feat(graph): add per-node timing and datastore-read attribution to check debug trace

Adds start_time and datastore_queries to the internal dispatch CheckDebugTrace
so a single Check / CheckBulk debug trace can answer two questions that
duration alone cannot: (1) whether sub-problems were traversed concurrently,
and (2) where time was spent blocked on datastore reads.

start_time is the wall-clock anchor for each node (the monotonic duration stays
authoritative; end = start_time + duration). A context-scoped, concurrency-safe
collector is installed per dispatch node; the observable datastore proxy records
each relationship query (shape, start, duration, row count) into it, so reads
attribute to the node that issued them. Both only activate when debug tracing is
enabled, leaving the non-debug hot path untouched.

This is the internal-proto-only slice; the public authzed/api CheckDebugTrace
fields and the internal->public conversion follow once that field lands.

Refs #3198

Author: poucet

internal/datastore/proxy/observable.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"context"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -12,6 +13,7 @@ import (
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/internal/dstrace"
 	"github.com/authzed/spicedb/internal/telemetry/otelconv"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
@@ -230,12 +232,19 @@ func (r *observableReader) LegacyReadNamespaceByName(ctx context.Context, nsName
 
 func (r *observableReader) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, opts ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) {
 	queryOpts := options.NewQueryOptionsWithOptions(opts...)
-	ctx, closer := observe(ctx, "QueryRelationships", string(queryOpts.QueryShape), trace.WithAttributes(
+	queryShape := string(queryOpts.QueryShape)
+	ctx, closer := observe(ctx, "QueryRelationships", queryShape, trace.WithAttributes(
 		attribute.String(otelconv.AttrDatastoreResourceType, filter.OptionalResourceType),
 		attribute.String(otelconv.AttrDatastoreResourceRelation, filter.OptionalResourceRelation),
-		attribute.String(otelconv.AttrDatastoreQueryShape, string(queryOpts.QueryShape)),
+		attribute.String(otelconv.AttrDatastoreQueryShape, queryShape),
 	))
 
+	collector := dstrace.CollectorFromContext(ctx)
+	var start time.Time
+	if collector != nil {
+		start = time.Now()
+	}
+
 	iterator, err := r.delegate.QueryRelationships(ctx, filter, opts...)
 	if err != nil {
 		closer()
@@ -253,14 +262,24 @@ func (r *observableReader) QueryRelationships(ctx context.Context, filter datast
 			}
 		}
 		loadedRelationshipCount.Observe(float64(count))
+		if collector != nil {
+			collector.Record(queryShape, start, time.Since(start), count)
+		}
 	}, nil
 }
 
 func (r *observableReader) ReverseQueryRelationships(ctx context.Context, subjectsFilter datastore.SubjectsFilter, opts ...options.ReverseQueryOptionsOption) (datastore.RelationshipIterator, error) {
 	queryOpts := options.NewReverseQueryOptionsWithOptions(opts...)
-	ctx, closer := observe(ctx, "ReverseQueryRelationships", string(queryOpts.QueryShapeForReverse), trace.WithAttributes(
+	queryShape := string(queryOpts.QueryShapeForReverse)
+	ctx, closer := observe(ctx, "ReverseQueryRelationships", queryShape, trace.WithAttributes(
 		attribute.String(otelconv.AttrDatastoreSubjectType, subjectsFilter.SubjectType),
-		attribute.String(otelconv.AttrDatastoreQueryShape, string(queryOpts.QueryShapeForReverse))))
+		attribute.String(otelconv.AttrDatastoreQueryShape, queryShape)))
+
+	collector := dstrace.CollectorFromContext(ctx)
+	var start time.Time
+	if collector != nil {
+		start = time.Now()
+	}
 
 	iterator, err := r.delegate.ReverseQueryRelationships(ctx, subjectsFilter, opts...)
 	if err != nil {
@@ -279,6 +298,9 @@ func (r *observableReader) ReverseQueryRelationships(ctx context.Context, subjec
 			}
 		}
 		loadedRelationshipCount.Observe(float64(count))
+		if collector != nil {
+			collector.Record(queryShape, start, time.Since(start), count)
+		}
 	}, nil
 }
 

internal/datastore/proxy/observable_test.go

@@ -13,7 +13,10 @@ import (
 
 	"github.com/authzed/spicedb/internal/datastore/proxy/proxy_test"
 	"github.com/authzed/spicedb/internal/datastore/revisions"
+	"github.com/authzed/spicedb/internal/dstrace"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
@@ -360,6 +363,90 @@ func TestObservableProxy_ReaderMethodsWithMetrics(t *testing.T) {
 	}
 }
 
+func TestObservableProxy_RecordsDatastoreQueriesToCollector(t *testing.T) {
+	t.Run("QueryRelationships", func(t *testing.T) {
+		dsMock, readerMock, _ := newMocks()
+		twoRelIter := datastore.RelationshipIterator(func(yield func(tuple.Relationship, error) bool) {
+			if !yield(tuple.MustParse("document:1#viewer@user:1"), nil) {
+				return
+			}
+			yield(tuple.MustParse("document:1#viewer@user:2"), nil)
+		})
+		readerMock.On("QueryRelationships", datastore.RelationshipsFilter{OptionalResourceType: "document"}, mock.Anything).
+			Return(twoRelIter, nil).Once()
+
+		sut := NewObservableDatastoreProxy(dsMock)
+		ctx, collector := dstrace.WithCollector(t.Context())
+
+		iter, err := sut.SnapshotReader(testRev).QueryRelationships(
+			ctx,
+			datastore.RelationshipsFilter{OptionalResourceType: "document"},
+			options.WithQueryShape(queryshape.CheckPermissionSelectDirectSubjects),
+		)
+		require.NoError(t, err)
+		for range iter {
+		}
+
+		queries := collector.Queries()
+		require.Len(t, queries, 1)
+		require.Equal(t, string(queryshape.CheckPermissionSelectDirectSubjects), queries[0].QueryShape)
+		require.Equal(t, uint64(2), queries[0].RelationshipCount)
+		require.NotNil(t, queries[0].StartTime)
+		require.NotNil(t, queries[0].Duration)
+
+		dsMock.AssertExpectations(t)
+		readerMock.AssertExpectations(t)
+	})
+
+	t.Run("ReverseQueryRelationships", func(t *testing.T) {
+		dsMock, readerMock, _ := newMocks()
+		oneRelIter := datastore.RelationshipIterator(func(yield func(tuple.Relationship, error) bool) {
+			yield(tuple.MustParse("document:1#viewer@user:1"), nil)
+		})
+		readerMock.On("ReverseQueryRelationships", datastore.SubjectsFilter{SubjectType: "user"}, mock.Anything).
+			Return(oneRelIter, nil).Once()
+
+		sut := NewObservableDatastoreProxy(dsMock)
+		ctx, collector := dstrace.WithCollector(t.Context())
+
+		iter, err := sut.SnapshotReader(testRev).ReverseQueryRelationships(
+			ctx,
+			datastore.SubjectsFilter{SubjectType: "user"},
+			options.WithQueryShapeForReverse(queryshape.MatchingResourcesForSubject),
+		)
+		require.NoError(t, err)
+		for range iter {
+		}
+
+		queries := collector.Queries()
+		require.Len(t, queries, 1)
+		require.Equal(t, string(queryshape.MatchingResourcesForSubject), queries[0].QueryShape)
+		require.Equal(t, uint64(1), queries[0].RelationshipCount)
+
+		dsMock.AssertExpectations(t)
+		readerMock.AssertExpectations(t)
+	})
+
+	t.Run("no collector is a no-op", func(t *testing.T) {
+		dsMock, readerMock, _ := newMocks()
+		iter := datastore.RelationshipIterator(func(yield func(tuple.Relationship, error) bool) {
+			yield(tuple.MustParse("document:1#viewer@user:1"), nil)
+		})
+		readerMock.On("QueryRelationships", datastore.RelationshipsFilter{OptionalResourceType: "document"}).
+			Return(iter, nil).Once()
+
+		sut := NewObservableDatastoreProxy(dsMock)
+		got, err := sut.SnapshotReader(testRev).QueryRelationships(
+			t.Context(), datastore.RelationshipsFilter{OptionalResourceType: "document"})
+		require.NoError(t, err)
+		for range got {
+		}
+
+		dsMock.AssertExpectations(t)
+		readerMock.AssertExpectations(t)
+	})
+}
+
 func TestObservableProxy_RWTMethodsWithMetrics(t *testing.T) {
 	tests := []struct {
 		name      string

internal/dstrace/collector.go

@@ -0,0 +1,71 @@
+// Package dstrace provides a context-scoped collector for attributing datastore
+// query timings to a single Check dispatch node. It is populated by the
+// observable datastore proxy and drained into the check debug trace, but lives
+// in its own leaf package so neither side has to depend on the other.
+package dstrace
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
+)
+
+type collectorCtxKey struct{}
+
+// Collector accumulates datastore query observations for a single Check
+// dispatch node. It is safe for concurrent use: a node may issue multiple
+// relationship queries in parallel (e.g. the branches of a union).
+type Collector struct {
+	mu      sync.Mutex
+	queries []*v1.DatastoreQuery
+}
+
+// WithCollector returns a context carrying a fresh Collector, along with that
+// collector. It should be called once per dispatch node when debug tracing is
+// enabled so that the queries the node issues itself are attributed to it;
+// dispatched children re-install their own collector and therefore capture their
+// own queries.
+func WithCollector(ctx context.Context) (context.Context, *Collector) {
+	c := &Collector{}
+	return context.WithValue(ctx, collectorCtxKey{}, c), c
+}
+
+// CollectorFromContext returns the Collector stored in ctx, or nil if none is
+// present (i.e. debug tracing is disabled). A nil Collector is safe to call
+// Record/Queries on.
+func CollectorFromContext(ctx context.Context) *Collector {
+	c, _ := ctx.Value(collectorCtxKey{}).(*Collector)
+	return c
+}
+
+// Record appends a single datastore query observation. It is a no-op on a nil
+// Collector.
+func (c *Collector) Record(queryShape string, start time.Time, duration time.Duration, relationshipCount uint64) {
+	if c == nil {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.queries = append(c.queries, &v1.DatastoreQuery{
+		QueryShape:        queryShape,
+		StartTime:         timestamppb.New(start),
+		Duration:          durationpb.New(duration),
+		RelationshipCount: relationshipCount,
+	})
+}
+
+// Queries returns the collected datastore query observations. It returns nil on
+// a nil Collector.
+func (c *Collector) Queries() []*v1.DatastoreQuery {
+	if c == nil {
+		return nil
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.queries
+}

internal/dstrace/collector_test.go

@@ -0,0 +1,78 @@
+package dstrace
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCollectorFromContextAbsentIsNilSafe(t *testing.T) {
+	c := CollectorFromContext(context.Background())
+	require.Nil(t, c)
+
+	// A nil collector must be safe to use.
+	c.Record("shape", time.Now(), time.Millisecond, 1)
+	require.Nil(t, c.Queries())
+}
+
+func TestWithCollectorRecordsQueries(t *testing.T) {
+	ctx, c := WithCollector(context.Background())
+	require.NotNil(t, c)
+	require.Same(t, c, CollectorFromContext(ctx))
+
+	start := time.Now()
+	c.Record("shape-a", start, 2*time.Millisecond, 3)
+	c.Record("shape-b", start, 5*time.Millisecond, 0)
+
+	queries := c.Queries()
+	require.Len(t, queries, 2)
+
+	require.Equal(t, "shape-a", queries[0].QueryShape)
+	require.Equal(t, uint64(3), queries[0].RelationshipCount)
+	require.Equal(t, 2*time.Millisecond, queries[0].Duration.AsDuration())
+	require.WithinDuration(t, start, queries[0].StartTime.AsTime(), time.Microsecond)
+
+	require.Equal(t, "shape-b", queries[1].QueryShape)
+	require.Equal(t, uint64(0), queries[1].RelationshipCount)
+}
+
+func TestCollectorConcurrentRecord(t *testing.T) {
+	_, c := WithCollector(context.Background())
+
+	const goroutines = 16
+	const perGoroutine = 64
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	for i := 0; i < goroutines; i++ {
+		go func() {
+			defer wg.Done()
+			for j := 0; j < perGoroutine; j++ {
+				c.Record("shape", time.Now(), time.Microsecond, 1)
+			}
+		}()
+	}
+	wg.Wait()
+
+	require.Len(t, c.Queries(), goroutines*perGoroutine)
+}
+
+// Each child context gets its own collector, so a parent never sees a child's
+// recorded queries.
+func TestNestedCollectorsAreIndependent(t *testing.T) {
+	parentCtx, parent := WithCollector(context.Background())
+	parent.Record("parent-query", time.Now(), time.Millisecond, 1)
+
+	childCtx, child := WithCollector(parentCtx)
+	child.Record("child-query", time.Now(), time.Millisecond, 2)
+
+	require.Len(t, parent.Queries(), 1)
+	require.Equal(t, "parent-query", parent.Queries()[0].QueryShape)
+
+	require.Same(t, child, CollectorFromContext(childCtx))
+	require.Len(t, child.Queries(), 1)
+	require.Equal(t, "child-query", child.Queries()[0].QueryShape)
+}

internal/graph/check.go

@@ -10,8 +10,10 @@ import (
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/authzed/spicedb/internal/dispatch"
+	"github.com/authzed/spicedb/internal/dstrace"
 	"github.com/authzed/spicedb/internal/graph/hints"
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/namespace"
@@ -98,9 +100,14 @@ type currentRequestContext struct {
 // Check performs a check request with the provided request and context
 func (cc *ConcurrentChecker) Check(ctx context.Context, req ValidatedCheckRequest, relation *core.Relation) (*v1.DispatchCheckResponse, error) {
 	var startTime *time.Time
+	var collector *dstrace.Collector
 	if req.Debug != v1.DispatchCheckRequest_NO_DEBUG {
 		now := time.Now()
 		startTime = &now
+		// Install a per-node collector so the datastore queries this node issues
+		// itself are attributed to it. Dispatched children re-enter Check and
+		// install their own collector, so their reads attribute to them.
+		ctx, collector = dstrace.WithCollector(ctx)
 	}
 
 	resolved := cc.checkInternal(ctx, req, relation)
@@ -131,6 +138,8 @@ func (cc *ConcurrentChecker) Check(ctx context.Context, req ValidatedCheckReques
 
 	debugInfo.Check.Request = clonedRequest
 	debugInfo.Check.Duration = durationpb.New(time.Since(*startTime))
+	debugInfo.Check.StartTime = timestamppb.New(*startTime)
+	debugInfo.Check.DatastoreQueries = collector.Queries()
 
 	if nspkg.GetRelationKind(relation) == iv1.RelationMetadata_PERMISSION {
 		debugInfo.Check.ResourceRelationType = v1.CheckDebugTrace_PERMISSION