Merge pull request #1595 from authzed/improve-tracing-ux

Some improvements to tracing UX

internal/datastore/common/helpers.go

@@ -25,7 +25,7 @@ func WriteTuples(ctx context.Context, ds datastore.Datastore, op core.RelationTu
 
 // UpdateTuplesInDatastore is a convenience method to perform multiple relation update operations on a Datastore
 func UpdateTuplesInDatastore(ctx context.Context, ds datastore.Datastore, updates ...*core.RelationTupleUpdate) (datastore.Revision, error) {
-	return ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	return ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteRelationships(ctx, updates)
 	})
 }

internal/datastore/common/revisions/optimized.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/benbjohnson/clock"
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/singleflight"
 
 	log "github.com/authzed/spicedb/internal/logging"
@@ -37,9 +38,7 @@ func (cor *CachedOptimizedRevisions) SetOptimizedRevisionFunc(revisionFunc Optim
 }
 
 func (cor *CachedOptimizedRevisions) OptimizedRevision(ctx context.Context) (datastore.Revision, error) {
-	ctx, span := tracer.Start(ctx, "OptimizedRevision")
-	defer span.End()
-
+	span := trace.SpanFromContext(ctx)
 	localNow := cor.clockFn.Now()
 
 	// Subtract a random amount of time from now, to let barely expired candidates get selected

internal/datastore/common/sql.go

@@ -496,8 +496,6 @@ func (tqs QueryExecutor) ExecuteQuery(
 	query SchemaQueryFilterer,
 	opts ...options.QueryOptionsOption,
 ) (datastore.RelationshipIterator, error) {
-	ctx, span := tracer.Start(ctx, "ExecuteQuery")
-	defer span.End()
 	queryOpts := options.NewQueryOptionsWithOptions(opts...)
 
 	query = query.TupleOrder(queryOpts.Sort)

internal/datastore/crdb/crdb.go

@@ -306,7 +306,7 @@ func (cds *crdbDatastore) ReadWriteTx(
 			0,
 		}
 
-		if err := f(rwt); err != nil {
+		if err := f(ctx, rwt); err != nil {
 			return err
 		}
 

internal/datastore/crdb/pool_test.go

@@ -121,7 +121,7 @@ func TestTxReset(t *testing.T) {
 
 			// WriteNamespace utilizes execute so we'll use it
 			i := 0
-			rev, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+			rev, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 				if i < len(tt.errors) {
 					defer func() { i++ }()
 					return tt.errors[i]

internal/datastore/memdb/memdb.go

@@ -131,7 +131,7 @@ func (mdb *memdbDatastore) SnapshotReader(revisionRaw datastore.Revision) datast
 }
 
 func (mdb *memdbDatastore) ReadWriteTx(
-	_ context.Context,
+	ctx context.Context,
 	f datastore.TxUserFunc,
 	opts ...options.RWTOptionsOption,
 ) (datastore.Revision, error) {
@@ -170,7 +170,7 @@ func (mdb *memdbDatastore) ReadWriteTx(
 
 		newRevision := mdb.newRevisionID()
 		rwt := &memdbReadWriteTx{memdbReader{&sync.Mutex{}, txSrc, nil}, newRevision}
-		if err := f(rwt); err != nil {
+		if err := f(ctx, rwt); err != nil {
 			mdb.Lock()
 			if tx != nil {
 				tx.Abort()

internal/datastore/memdb/memdb_test.go

@@ -46,7 +46,7 @@ func TestConcurrentWritePanic(t *testing.T) {
 
 	numPanics := uint64(0)
 	require.Eventually(func() bool {
-		_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			g := errgroup.Group{}
 			g.Go(func() (err error) {
 				defer func() {
@@ -94,7 +94,7 @@ func TestConcurrentWriteRelsError(t *testing.T) {
 	for i := 0; i < 50; i++ {
 		i := i
 		g.Go(func() error {
-			_, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+			_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 				updates := []*corev1.RelationTupleUpdate{}
 				for j := 0; j < 500; j++ {
 					updates = append(updates, &corev1.RelationTupleUpdate{

internal/datastore/mysql/datastore.go

@@ -314,7 +314,7 @@ func (mds *Datastore) ReadWriteTx(
 				newTxnID,
 			}
 
-			return fn(rwt)
+			return fn(ctx, rwt)
 		}); err != nil {
 			if !config.DisableRetries && isErrorRetryable(err) {
 				continue

internal/datastore/mysql/datastore_test.go

@@ -181,7 +181,7 @@ func GarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	req.True(r.IsReady)
 
 	// Write basic namespaces.
-	writtenAt, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	writtenAt, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(
@@ -202,7 +202,7 @@ func GarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	req.Zero(removed.Namespaces)
 
 	// Replace the namespace with a new one.
-	writtenAt, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	writtenAt, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(
@@ -322,7 +322,7 @@ func GarbageCollectionByTimeTest(t *testing.T, ds datastore.Datastore) {
 	req.True(r.IsReady)
 
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(
@@ -419,7 +419,7 @@ func NoRelationshipsGarbageCollectionTest(t *testing.T, ds datastore.Datastore)
 	req.True(r.IsReady)
 
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(
@@ -456,7 +456,7 @@ func ChunkedGarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	req.True(r.IsReady)
 
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(

internal/datastore/postgres/postgres.go

@@ -347,7 +347,7 @@ func (pgd *pgDatastore) ReadWriteTx(
 				newXID,
 			}
 
-			return fn(rwt)
+			return fn(ctx, rwt)
 		}))
 
 		if err != nil {

internal/datastore/postgres/postgres_test.go

@@ -239,7 +239,7 @@ func GarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	r, err := ds.ReadyState(ctx)
 	require.NoError(err)
 	require.True(r.IsReady)
-	firstWrite, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	firstWrite, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		// Write basic namespaces.
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
@@ -258,7 +258,7 @@ func GarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	require.Zero(removed.Namespaces)
 
 	// Replace the namespace with a new one.
-	updateTwoNamespaces, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	updateTwoNamespaces, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(
 			ctx,
 			namespace.Namespace(
@@ -367,7 +367,7 @@ func GarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	tRequire.TupleExists(ctx, tpl, relLastWriteAt)
 
 	// Inject a transaction to clean up the last write
-	lastRev, err := pds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	lastRev, err := pds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return nil
 	})
 	require.NoError(err)
@@ -426,7 +426,7 @@ func GarbageCollectionByTimeTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 	require.True(r.IsReady)
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
 			namespace.MustRelation("reader", nil),
@@ -469,7 +469,7 @@ func GarbageCollectionByTimeTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 
 	// Inject a revision to sweep up the last revision
-	_, err = pds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = pds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return nil
 	})
 	require.NoError(err)
@@ -501,7 +501,7 @@ func ChunkedGarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 	require.True(r.IsReady)
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
 			namespace.MustRelation("reader", nil),
@@ -549,7 +549,7 @@ func ChunkedGarbageCollectionTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 
 	// Inject a revision to sweep up the last revision
-	_, err = pds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = pds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return nil
 	})
 	require.NoError(err)
@@ -703,7 +703,7 @@ func ConcurrentRevisionHeadTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 	require.True(r.IsReady)
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
 			namespace.MustRelation("reader", nil),
@@ -719,7 +719,7 @@ func ConcurrentRevisionHeadTest(t *testing.T, ds datastore.Datastore) {
 	var commitLastRev, commitFirstRev datastore.Revision
 	g.Go(func() error {
 		var err error
-		commitLastRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		commitLastRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			rtu := tuple.Touch(&core.RelationTuple{
 				ResourceAndRelation: &core.ObjectAndRelation{
 					Namespace: "resource",
@@ -746,7 +746,7 @@ func ConcurrentRevisionHeadTest(t *testing.T, ds datastore.Datastore) {
 
 	<-waitToStart
 
-	commitFirstRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	commitFirstRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		rtu := tuple.Touch(&core.RelationTuple{
 			ResourceAndRelation: &core.ObjectAndRelation{
 				Namespace: "resource",
@@ -805,7 +805,7 @@ func ConcurrentRevisionWatchTest(t *testing.T, ds datastore.Datastore) {
 	require.True(r.IsReady)
 
 	// Write basic namespaces.
-	rev, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	rev, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
 			namespace.MustRelation("reader", nil),
@@ -852,7 +852,7 @@ func ConcurrentRevisionWatchTest(t *testing.T, ds datastore.Datastore) {
 	var commitLastRev, commitFirstRev datastore.Revision
 	g.Go(func() error {
 		var err error
-		commitLastRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		commitLastRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			err = rwt.WriteRelationships(ctx, []*core.RelationTupleUpdate{
 				tuple.Touch(tuple.MustParse("something:001#viewer@user:123")),
 				tuple.Touch(tuple.MustParse("something:002#viewer@user:123")),
@@ -871,7 +871,7 @@ func ConcurrentRevisionWatchTest(t *testing.T, ds datastore.Datastore) {
 
 	<-waitToStart
 
-	commitFirstRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	commitFirstRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteRelationships(ctx, []*core.RelationTupleUpdate{
 			tuple.Touch(tuple.MustParse("resource:1001#reader@user:456")),
 			tuple.Touch(tuple.MustParse("resource:1002#reader@user:456")),
@@ -889,7 +889,7 @@ func ConcurrentRevisionWatchTest(t *testing.T, ds datastore.Datastore) {
 	require.False(commitFirstRev.Equal(commitLastRev))
 
 	// Write another revision.
-	afterRev, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	afterRev, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		rtu := tuple.Touch(&core.RelationTuple{
 			ResourceAndRelation: &core.ObjectAndRelation{
 				Namespace: "resource",
@@ -1031,7 +1031,7 @@ func RevisionInversionTest(t *testing.T, ds datastore.Datastore) {
 	require.NoError(err)
 	require.True(r.IsReady)
 	// Write basic namespaces.
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace(
 			"resource",
 			namespace.MustRelation("reader", nil),
@@ -1047,7 +1047,7 @@ func RevisionInversionTest(t *testing.T, ds datastore.Datastore) {
 	var commitLastRev, commitFirstRev datastore.Revision
 	g.Go(func() error {
 		var err error
-		commitLastRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		commitLastRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			rtu := tuple.Touch(&core.RelationTuple{
 				ResourceAndRelation: &core.ObjectAndRelation{
 					Namespace: "resource",
@@ -1074,7 +1074,7 @@ func RevisionInversionTest(t *testing.T, ds datastore.Datastore) {
 
 	<-waitToStart
 
-	commitFirstRev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	commitFirstRev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		rtu := tuple.Touch(&core.RelationTuple{
 			ResourceAndRelation: &core.ObjectAndRelation{
 				Namespace: "resource",
@@ -1112,7 +1112,7 @@ func OTelTracingTest(t *testing.T, ds datastore.Datastore) {
 	testTraceProvider.RegisterSpanProcessor(spanrecorder)
 
 	// Perform basic operation
-	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, namespace.Namespace("resource"))
 	})
 	require.NoError(err)
@@ -1207,7 +1207,7 @@ func datastoreWithInterceptorAndTestData(t *testing.T, interceptor pgcommon.Quer
 	// Write namespaces and a few thousand relationships.
 	ctx := context.Background()
 	for i := 0; i < 1000; i++ {
-		_, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			err := rwt.WriteNamespaces(ctx, namespace.Namespace(
 				fmt.Sprintf("resource%d", i),
 				namespace.MustRelation("reader", nil)))
@@ -1262,7 +1262,7 @@ func datastoreWithInterceptorAndTestData(t *testing.T, interceptor pgcommon.Quer
 
 	// Delete some relationships.
 	for i := 990; i < 1000; i++ {
-		_, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			rtu := tuple.Delete(&core.RelationTuple{
 				ResourceAndRelation: &core.ObjectAndRelation{
 					Namespace: testfixtures.DocumentNS.Name,
@@ -1295,7 +1295,7 @@ func datastoreWithInterceptorAndTestData(t *testing.T, interceptor pgcommon.Quer
 
 	// Write some more relationships.
 	for i := 1000; i < 1100; i++ {
-		_, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 			// Write some relationships.
 			rtu := tuple.Touch(&core.RelationTuple{
 				ResourceAndRelation: &core.ObjectAndRelation{

internal/datastore/proxy/observable.go

@@ -77,8 +77,8 @@ func (p *observableProxy) ReadWriteTx(
 	f datastore.TxUserFunc,
 	opts ...options.RWTOptionsOption,
 ) (datastore.Revision, error) {
-	return p.delegate.ReadWriteTx(ctx, func(delegateRWT datastore.ReadWriteTransaction) error {
-		return f(&observableRWT{&observableReader{delegateRWT}, delegateRWT})
+	return p.delegate.ReadWriteTx(ctx, func(ctx context.Context, delegateRWT datastore.ReadWriteTransaction) error {
+		return f(ctx, &observableRWT{&observableReader{delegateRWT}, delegateRWT})
 	}, opts...)
 }
 
@@ -193,7 +193,11 @@ func (r *observableReader) ReadNamespaceByName(ctx context.Context, nsName strin
 }
 
 func (r *observableReader) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, options ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) {
-	ctx, closer := observe(ctx, "QueryRelationships")
+	ctx, closer := observe(ctx, "QueryRelationships", trace.WithAttributes(
+		attribute.String("resourceType", filter.ResourceType),
+		attribute.String("resourceRelation", filter.OptionalResourceRelation),
+		attribute.String("caveatName", filter.OptionalCaveatName),
+	))
 
 	iterator, err := r.delegate.QueryRelationships(ctx, filter, options...)
 	if err != nil {

internal/datastore/proxy/proxy_test/mock.go

@@ -21,14 +21,14 @@ func (dm *MockDatastore) SnapshotReader(rev datastore.Revision) datastore.Reader
 }
 
 func (dm *MockDatastore) ReadWriteTx(
-	_ context.Context,
+	ctx context.Context,
 	f datastore.TxUserFunc,
 	opts ...options.RWTOptionsOption,
 ) (datastore.Revision, error) {
 	args := dm.Called(opts)
 	mockRWT := args.Get(0).(datastore.ReadWriteTransaction)
 
-	if err := f(mockRWT); err != nil {
+	if err := f(ctx, mockRWT); err != nil {
 		return datastore.NoRevision, err
 	}
 

internal/datastore/proxy/readonly_test.go

@@ -34,13 +34,13 @@ func TestRWOperationErrors(t *testing.T) {
 	ds := NewReadonlyDatastore(delegate)
 	ctx := context.Background()
 
-	rev, err := ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	rev, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.DeleteNamespaces(ctx, "fake")
 	})
 	require.ErrorAs(err, &datastore.ErrReadOnly{})
 	require.Equal(datastore.NoRevision, rev)
 
-	rev, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+	rev, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.WriteNamespaces(ctx, &core.NamespaceDefinition{Name: "user"})
 	})
 	require.ErrorAs(err, &datastore.ErrReadOnly{})