refactor: remove analyze.go, it is now observer_analyze.go and simplify

Author: barakmich

pkg/query/analyze.go

@@ -2,7 +2,6 @@ package query
 
 import (
 	"context"
-	"time"
 
 	"github.com/authzed/spicedb/internal/caveats"
 	"github.com/authzed/spicedb/pkg/datalayer"
@@ -184,18 +183,15 @@ func (ctx *Context) notifyEnterIterator(op ObserverOperation, key CanonicalKey)
 	}
 }
 
-// wrapPathSeqWithObservers wraps a PathSeq to notify all registered observers of paths and timing.
+// wrapPathSeqWithObservers wraps a PathSeq to notify all registered observers of paths and completion.
 // Returns the original PathSeq unchanged when there are no observers.
 func (ctx *Context) wrapPathSeqWithObservers(op ObserverOperation, key CanonicalKey, pathSeq PathSeq) PathSeq {
 	if len(ctx.observers) == 0 {
 		return pathSeq
 	}
 	return func(yield func(Path, error) bool) {
-		start := time.Now()
 		defer func() {
-			dur := time.Since(start)
 			for _, obs := range ctx.observers {
-				obs.ObserveTiming(op, key, dur)
 				obs.ObserveReturnIterator(op, key)
 			}
 		}()

pkg/query/context.go

@@ -1,7 +1,5 @@
 package query
 
-import "time"
-
 type ObserverOperation int
 
 const (
@@ -11,8 +9,7 @@ const (
 )
 
 type Observer interface {
-	ObserveTiming(op ObserverOperation, key CanonicalKey, dur time.Duration)
-	ObservePath(op ObserverOperation, key CanonicalKey, path Path)
 	ObserveEnterIterator(op ObserverOperation, key CanonicalKey)
+	ObservePath(op ObserverOperation, key CanonicalKey, path Path)
 	ObserveReturnIterator(op ObserverOperation, key CanonicalKey)
 }

pkg/query/observer.go

@@ -1,7 +1,9 @@
 package query
 
 import (
+	"fmt"
 	"maps"
+	"strings"
 	"sync"
 	"time"
 )
@@ -21,18 +23,20 @@ type AnalyzeStats struct {
 
 // AnalyzeObserver is a thread-safe Observer that collects execution statistics keyed by CanonicalKey.
 type AnalyzeObserver struct {
-	mu    sync.Mutex
-	stats map[CanonicalKey]AnalyzeStats // GUARDED_BY(mu)
+	mu         sync.Mutex
+	stats      map[CanonicalKey]AnalyzeStats                    // GUARDED_BY(mu)
+	startTimes map[ObserverOperation]map[CanonicalKey]time.Time // GUARDED_BY(mu)
 }
 
 // NewAnalyzeObserver creates a new thread-safe analyze observer.
 func NewAnalyzeObserver() *AnalyzeObserver {
 	return &AnalyzeObserver{
-		stats: make(map[CanonicalKey]AnalyzeStats),
+		stats:      make(map[CanonicalKey]AnalyzeStats),
+		startTimes: make(map[ObserverOperation]map[CanonicalKey]time.Time),
 	}
 }
 
-// ObserveEnterIterator increments the call counter for the given iterator key and operation.
+// ObserveEnterIterator increments the call counter and records the start time.
 func (a *AnalyzeObserver) ObserveEnterIterator(op ObserverOperation, key CanonicalKey) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
@@ -47,6 +51,11 @@ func (a *AnalyzeObserver) ObserveEnterIterator(op ObserverOperation, key Canonic
 		stats.IterResourcesCalls++
 	}
 	a.stats[key] = stats
+
+	if a.startTimes[op] == nil {
+		a.startTimes[op] = make(map[CanonicalKey]time.Time)
+	}
+	a.startTimes[op][key] = time.Now()
 }
 
 // ObservePath increments the result counter for the given iterator key and operation.
@@ -66,26 +75,29 @@ func (a *AnalyzeObserver) ObservePath(op ObserverOperation, key CanonicalKey, _
 	a.stats[key] = stats
 }
 
-// ObserveTiming records the elapsed time for the given iterator key and operation.
-func (a *AnalyzeObserver) ObserveTiming(op ObserverOperation, key CanonicalKey, dur time.Duration) {
+// ObserveReturnIterator records elapsed time since ObserveEnterIterator was called.
+func (a *AnalyzeObserver) ObserveReturnIterator(op ObserverOperation, key CanonicalKey) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	stats := a.stats[key]
-	switch op {
-	case CheckOperation:
-		stats.CheckTime += dur
-	case IterSubjectsOperation:
-		stats.IterSubjectsTime += dur
-	case IterResourcesOperation:
-		stats.IterResourcesTime += dur
+	if opTimes := a.startTimes[op]; opTimes != nil {
+		if start, ok := opTimes[key]; ok {
+			elapsed := time.Since(start)
+			stats := a.stats[key]
+			switch op {
+			case CheckOperation:
+				stats.CheckTime += elapsed
+			case IterSubjectsOperation:
+				stats.IterSubjectsTime += elapsed
+			case IterResourcesOperation:
+				stats.IterResourcesTime += elapsed
+			}
+			a.stats[key] = stats
+			delete(opTimes, key)
+		}
 	}
-	a.stats[key] = stats
 }
 
-// ObserveReturnIterator is a no-op for the analyze observer.
-func (a *AnalyzeObserver) ObserveReturnIterator(_ ObserverOperation, _ CanonicalKey) {}
-
 // GetStats returns a copy of all collected stats.
 func (a *AnalyzeObserver) GetStats() map[CanonicalKey]AnalyzeStats {
 	a.mu.Lock()
@@ -95,3 +107,107 @@ func (a *AnalyzeObserver) GetStats() map[CanonicalKey]AnalyzeStats {
 	maps.Copy(result, a.stats)
 	return result
 }
+
+// FormatAnalysis returns a formatted string showing the iterator tree with execution statistics
+// for each iterator. Stats are looked up by iterator canonical key from the analyze map.
+func FormatAnalysis(tree Iterator, analyze map[CanonicalKey]AnalyzeStats) string {
+	if tree == nil {
+		return "No iterator tree provided"
+	}
+	if len(analyze) == 0 {
+		return "No analysis data available"
+	}
+
+	var sb strings.Builder
+	formatNode(tree, analyze, &sb, "", true)
+	return sb.String()
+}
+
+// AggregateAnalyzeStats combines all the analyze stats from a map into a single
+// aggregated AnalyzeStats. This is useful for getting total counts across all
+// iterators in a query execution.
+func AggregateAnalyzeStats(analyze map[CanonicalKey]AnalyzeStats) AnalyzeStats {
+	var total AnalyzeStats
+	for _, stats := range analyze {
+		total.CheckCalls += stats.CheckCalls
+		total.IterSubjectsCalls += stats.IterSubjectsCalls
+		total.IterResourcesCalls += stats.IterResourcesCalls
+		total.CheckResults += stats.CheckResults
+		total.IterSubjectsResults += stats.IterSubjectsResults
+		total.IterResourcesResults += stats.IterResourcesResults
+		total.CheckTime += stats.CheckTime
+		total.IterSubjectsTime += stats.IterSubjectsTime
+		total.IterResourcesTime += stats.IterResourcesTime
+	}
+	return total
+}
+
+// formatNode recursively formats a single iterator node and its children
+// depth parameter tracks how deep we are in the tree (0 = root)
+func formatNode(it Iterator, analyze map[CanonicalKey]AnalyzeStats, sb *strings.Builder, indent string, isLast bool) {
+	if it == nil {
+		return
+	}
+
+	// Get iterator info from Explain()
+	explain := it.Explain()
+	iterKey := it.CanonicalKey()
+	stats := analyze[iterKey]
+
+	// Draw tree branch for non-root nodes
+	// Root nodes have both empty indent and are drawn without any prefix
+	// Check if we're a child by seeing if we have any indentation OR if this is being called recursively
+	isRoot := indent == "" && !strings.Contains(sb.String(), "Calls:")
+
+	if !isRoot {
+		if isLast {
+			sb.WriteString(indent + "└─ ")
+		} else {
+			sb.WriteString(indent + "├─ ")
+		}
+	}
+
+	// Write iterator name and key (truncated for readability)
+	fmt.Fprintf(sb, "%s (Hash: %016x)\n", explain.Info, iterKey.Hash())
+
+	// Write stats with indentation
+	statsIndent := indent
+	if indent != "" {
+		if isLast {
+			statsIndent += "   "
+		} else {
+			statsIndent += "│  "
+		}
+	}
+
+	// Format call counts
+	sb.WriteString(statsIndent + fmt.Sprintf("   Calls: Check=%d, IterSubjects=%d, IterResources=%d\n",
+		stats.CheckCalls, stats.IterSubjectsCalls, stats.IterResourcesCalls))
+
+	// Format result counts
+	sb.WriteString(statsIndent + fmt.Sprintf("   Results: Check=%d, IterSubjects=%d, IterResources=%d\n",
+		stats.CheckResults, stats.IterSubjectsResults, stats.IterResourcesResults))
+
+	// Format timing
+	sb.WriteString(statsIndent + fmt.Sprintf("   Time: Check=%v, IterSubjects=%v, IterResources=%v\n",
+		stats.CheckTime, stats.IterSubjectsTime, stats.IterResourcesTime))
+
+	// Recursively format children
+	subIts := it.Subiterators()
+	if len(subIts) > 0 {
+		var childIndent string
+		switch {
+		case indent == "":
+			// Root node - children get no indent prefix, but will show tree chars
+			childIndent = ""
+		case isLast:
+			childIndent = indent + "   "
+		default:
+			childIndent = indent + "│  "
+		}
+
+		for i, subIt := range subIts {
+			formatNode(subIt, analyze, sb, childIndent, i == len(subIts)-1)
+		}
+	}
+}