fix: avoid deadlocking singleflighted reqs

Author: miparnisari

internal/datastore/proxy/singleflight.go

@@ -43,19 +43,29 @@ 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 {
 	_, _, err := p.checkRevGroup.Do(ctx, revision.String(), func(ctx context.Context) (string, error) {
-		return "", p.delegate.CheckRevision(ctx, revision)
+		// Sever the context so that a single caller's cancellation does not
+		// abort the query for all other singleflight waiters.
+		ctx, span := tracer.Start(ctx, "singleflightProxy.CheckRevision(sf)")
+		defer span.End()
+		return "", p.delegate.CheckRevision(context.WithoutCancel(ctx), revision)
 	})
 	return err
 }
 
 func (p *singleflightProxy) HeadRevision(ctx context.Context) (datastore.Revision, error) {
 	rev, _, err := p.headRevGroup.Do(ctx, "", func(ctx context.Context) (datastore.Revision, error) {
-		return p.delegate.HeadRevision(ctx)
+		// Sever the context so that a single caller's cancellation does not
+		// abort the query for all other singleflight waiters.
+		ctx, span := tracer.Start(ctx, "singleflightProxy.HeadRevision(sf)")
+		defer span.End()
+		return p.delegate.HeadRevision(context.WithoutCancel(ctx))
 	})
 	return rev, err
 }

internal/datastore/revisions/optimized.go

@@ -64,9 +64,14 @@ func (cor *CachedOptimizedRevisions) OptimizedRevision(ctx context.Context) (dat
 	cor.RUnlock()
 
 	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)
+		// Sever the context so that a single caller's cancellation does not
+		// abort the query for all other singleflight waiters. This is safe
+		// because optimizedFunc is a cheap read-only query (SELECT now()).
+		optimized, validFor, err := cor.optimizedFunc(context.WithoutCancel(ctx))
 		if err != nil {
 			return datastore.NoRevision, fmt.Errorf("unable to compute optimized revision: %w", err)
 		}

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,25 @@ func (p *ctxProxy) IsStrictReadModeEnabled() bool {
 }
 
 func (p *ctxProxy) OptimizedRevision(ctx context.Context) (Revision, error) {
-	return p.delegate.OptimizedRevision(context.WithoutCancel(ctx))
+	// NOTE: do NOT strip cancellation here. OptimizedRevision is singleflighted
+	// internally by CachedOptimizedRevisions, and stripping cancel before the
+	// singleflight makes it impossible for callers to escape the wait when the
+	// underlying query is slow. context.WithoutCancel is applied inside the
+	// singleflight function instead (see optimized.go).
+	ctx, span := tracer.Start(ctx, "ctxProxy.OptimizedRevision")
+	defer span.End()
+	return p.delegate.OptimizedRevision(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))
 }