fix: avoid deadlocking singleflighted reqs

Author: miparnisari

internal/datastore/proxy/singleflight.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"context"
+	"time"
 
 	"resenje.org/singleflight"
 
@@ -43,18 +44,30 @@ func (p *singleflightProxy) ReadWriteTx(ctx context.Context, f datastore.TxUserF
 func (p *singleflightProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, error) {
 	// NOTE: Optimized revisions are singleflighted by the underlying datastore via the
 	// CachedOptimizedRevisions struct.
+	ctx, span := tracer.Start(ctx, "singleflightProxy.OptimizedRevision")
+	defer span.End()
 	return p.delegate.OptimizedRevision(ctx)
 }
 
 func (p *singleflightProxy) CheckRevision(ctx context.Context, revision datastore.Revision) error {
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	_, _, err := p.checkRevGroup.Do(ctx, revision.String(), func(ctx context.Context) (string, error) {
+		ctx, span := tracer.Start(ctx, "singleflightProxy.CheckRevision(sf)")
+		defer span.End()
 		return "", p.delegate.CheckRevision(ctx, revision)
 	})
 	return err
 }
 
 func (p *singleflightProxy) HeadRevision(ctx context.Context) (datastore.Revision, error) {
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	rev, _, err := p.headRevGroup.Do(ctx, "", func(ctx context.Context) (datastore.Revision, error) {
+		ctx, span := tracer.Start(ctx, "singleflightProxy.HeadRevision(sf)")
+		defer span.End()
 		return p.delegate.HeadRevision(ctx)
 	})
 	return rev, err
@@ -69,6 +82,9 @@ func (p *singleflightProxy) Watch(ctx context.Context, afterRevision datastore.R
 }
 
 func (p *singleflightProxy) Statistics(ctx context.Context) (datastore.Stats, error) {
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	stats, _, err := p.statsGroup.Do(ctx, "", func(ctx context.Context) (datastore.Stats, error) {
 		return p.delegate.Statistics(ctx)
 	})

internal/datastore/revisions/optimized.go

@@ -63,7 +63,12 @@ func (cor *CachedOptimizedRevisions) OptimizedRevision(ctx context.Context) (dat
 	}
 	cor.RUnlock()
 
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	newQuantizedRevision, _, err := cor.updateGroup.Do(ctx, "", func(ctx context.Context) (datastore.Revision, error) {
+		ctx, span := tracer.Start(ctx, "CachedOptimizedRevisions.OptimizedRevision(sf)")
+		defer span.End()
 		log.Ctx(ctx).Debug().Time("now", localNow).Msg("computing new revision")
 
 		optimized, validFor, err := cor.optimizedFunc(ctx)

pkg/datastore/context.go

@@ -3,10 +3,14 @@ package datastore
 import (
 	"context"
 
+	"go.opentelemetry.io/otel"
+
 	"github.com/authzed/spicedb/pkg/datastore/options"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
 
+var tracer = otel.Tracer("spicedb/pkg/datastore/context")
+
 // NewSeparatingContextDatastoreProxy severs any timeouts in the context being
 // passed to the datastore and only retains tracing metadata.
 //
@@ -49,14 +53,20 @@ func (p *ctxProxy) IsStrictReadModeEnabled() bool {
 }
 
 func (p *ctxProxy) OptimizedRevision(ctx context.Context) (Revision, error) {
+	ctx, span := tracer.Start(ctx, "ctxProxy.OptimizedRevision")
+	defer span.End()
 	return p.delegate.OptimizedRevision(context.WithoutCancel(ctx))
 }
 
 func (p *ctxProxy) CheckRevision(ctx context.Context, revision Revision) error {
+	ctx, span := tracer.Start(ctx, "ctxProxy.CheckRevision")
+	defer span.End()
 	return p.delegate.CheckRevision(context.WithoutCancel(ctx), revision)
 }
 
 func (p *ctxProxy) HeadRevision(ctx context.Context) (Revision, error) {
+	ctx, span := tracer.Start(ctx, "ctxProxy.HeadRevision")
+	defer span.End()
 	return p.delegate.HeadRevision(context.WithoutCancel(ctx))
 }