feat(graph): per-node timing + datastore-read attribution in the check debug trace (internal proto)

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