feat(tracing): improved tracing for persisters and requests

internal/driver/daemon.go

@@ -11,7 +11,6 @@ import (
 
 	rts "github.com/ory/keto/proto/ory/keto/relation_tuples/v1alpha2"
 
-	"github.com/ory/x/otelx"
 	prometheus "github.com/ory/x/prometheusx"
 	grpcOtel "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"go.opentelemetry.io/otel"
@@ -36,6 +35,7 @@ import (
 	"github.com/ory/analytics-go/v4"
 	"github.com/ory/x/healthx"
 	"github.com/ory/x/metricsx"
+	"github.com/ory/x/otelx"
 	"github.com/spf13/cobra"
 
 	"github.com/ory/keto/internal/driver/config"
@@ -129,7 +129,7 @@ func (r *RegistryDefault) serveRead(ctx context.Context, done chan<- struct{}) f
 	rt, s := r.ReadRouter(ctx), r.ReadGRPCServer(ctx)
 
 	if tracer := r.Tracer(ctx); tracer.IsLoaded() {
-		rt = otelx.NewHandler(rt, "internal.driver.serveRead")
+		rt = otelx.TraceHandler(rt)
 	}
 
 	return func() error {
@@ -141,7 +141,7 @@ func (r *RegistryDefault) serveWrite(ctx context.Context, done chan<- struct{})
 	rt, s := r.WriteRouter(ctx), r.WriteGRPCServer(ctx)
 
 	if tracer := r.Tracer(ctx); tracer.IsLoaded() {
-		rt = otelx.NewHandler(rt, "internal.driver.serveWrite")
+		rt = otelx.TraceHandler(rt)
 	}
 
 	return func() error {

internal/persistence/sql/namespaces.go

@@ -7,6 +7,9 @@ import (
 )
 
 func (p *Persister) GetNamespaceByName(ctx context.Context, name string) (*namespace.Namespace, error) {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.GetNamespaceByName")
+	defer span.End()
+
 	nm, err := p.d.Config(ctx).NamespaceManager()
 	if err != nil {
 		return nil, err
@@ -15,6 +18,9 @@ func (p *Persister) GetNamespaceByName(ctx context.Context, name string) (*names
 }
 
 func (p *Persister) GetNamespaceByID(ctx context.Context, id int32) (*namespace.Namespace, error) {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.GetNamespaceByID")
+	defer span.End()
+
 	nm, err := p.d.Config(ctx).NamespaceManager()
 	if err != nil {
 		return nil, err

internal/persistence/sql/persister.go

@@ -34,6 +34,7 @@ type (
 	dependencies interface {
 		config.Provider
 		x.LoggerProvider
+		x.TracingProvider
 		ketoctx.ContextualizerProvider
 
 		PopConnection(ctx context.Context) (*pop.Connection, error)
@@ -75,6 +76,9 @@ func (p *Persister) Connection(ctx context.Context) *pop.Connection {
 }
 
 func (p *Persister) CreateWithNetwork(ctx context.Context, v interface{}) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.CreateWithNetwork")
+	defer span.End()
+
 	rv := reflect.ValueOf(v)
 
 	if rv.Kind() != reflect.Ptr && rv.Elem().Kind() != reflect.Struct {

internal/persistence/sql/relationtuples.go

@@ -113,6 +113,9 @@ func (r *RelationTuple) insertSubject(ctx context.Context, p *Persister, s relat
 }
 
 func (r *RelationTuple) FromInternal(ctx context.Context, p *Persister, rt *relationtuple.InternalRelationTuple) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.FromInternal")
+	defer span.End()
+
 	n, err := p.GetNamespaceByName(ctx, rt.Namespace)
 	if err != nil {
 		return err
@@ -126,6 +129,9 @@ func (r *RelationTuple) FromInternal(ctx context.Context, p *Persister, rt *rela
 }
 
 func (p *Persister) InsertRelationTuple(ctx context.Context, rel *relationtuple.InternalRelationTuple) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.InsertRelationTuple")
+	defer span.End()
+
 	if rel.Subject == nil {
 		return errors.WithStack(relationtuple.ErrNilSubject)
 	}
@@ -198,6 +204,9 @@ func (p *Persister) whereQuery(ctx context.Context, q *pop.Query, rq *relationtu
 }
 
 func (p *Persister) DeleteRelationTuples(ctx context.Context, rs ...*relationtuple.InternalRelationTuple) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.DeleteRelationTuples")
+	defer span.End()
+
 	return p.Transaction(ctx, func(ctx context.Context, _ *pop.Connection) error {
 		for _, r := range rs {
 			n, err := p.GetNamespaceByName(ctx, r.Namespace)
@@ -223,6 +232,9 @@ func (p *Persister) DeleteRelationTuples(ctx context.Context, rs ...*relationtup
 }
 
 func (p *Persister) DeleteAllRelationTuples(ctx context.Context, query *relationtuple.RelationQuery) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.DeleteAllRelationTuples")
+	defer span.End()
+
 	return p.Transaction(ctx, func(ctx context.Context, _ *pop.Connection) error {
 		sqlQuery := p.QueryWithNetwork(ctx)
 		err := p.whereQuery(ctx, sqlQuery, query)
@@ -236,6 +248,9 @@ func (p *Persister) DeleteAllRelationTuples(ctx context.Context, query *relation
 }
 
 func (p *Persister) GetRelationTuples(ctx context.Context, query *relationtuple.RelationQuery, options ...x.PaginationOptionSetter) ([]*relationtuple.InternalRelationTuple, string, error) {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.GetRelationTuples")
+	defer span.End()
+
 	pagination, err := internalPaginationFromOptions(options...)
 	if err != nil {
 		return nil, "", err
@@ -277,6 +292,9 @@ func (p *Persister) GetRelationTuples(ctx context.Context, query *relationtuple.
 }
 
 func (p *Persister) WriteRelationTuples(ctx context.Context, rs ...*relationtuple.InternalRelationTuple) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.WriteRelationTuples")
+	defer span.End()
+
 	return p.Transaction(ctx, func(ctx context.Context, _ *pop.Connection) error {
 		for _, r := range rs {
 			if err := p.InsertRelationTuple(ctx, r); err != nil {
@@ -288,6 +306,9 @@ func (p *Persister) WriteRelationTuples(ctx context.Context, rs ...*relationtupl
 }
 
 func (p *Persister) TransactRelationTuples(ctx context.Context, ins []*relationtuple.InternalRelationTuple, del []*relationtuple.InternalRelationTuple) error {
+	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.TransactRelationTuples")
+	defer span.End()
+
 	return p.Transaction(ctx, func(ctx context.Context, _ *pop.Connection) error {
 		if err := p.WriteRelationTuples(ctx, ins...); err != nil {
 			return err

internal/x/registry.go

@@ -1,8 +1,11 @@
 package x
 
 import (
+	"context"
+
 	"github.com/ory/herodot"
 	"github.com/ory/x/logrusx"
+	"github.com/ory/x/otelx"
 )
 
 type LoggerProvider interface {
@@ -12,3 +15,7 @@ type LoggerProvider interface {
 type WriterProvider interface {
 	Writer() herodot.Writer
 }
+
+type TracingProvider interface {
+	Tracer(ctx context.Context) *otelx.Tracer
+}