fix: put a timeout to callers of methods that use singleflight

Author: miparnisari

internal/datastore/crdb/crdb.go

@@ -20,7 +20,6 @@ import (
 	"github.com/shopspring/decimal"
 	"go.opentelemetry.io/otel"
 	"golang.org/x/sync/errgroup"
-	"resenje.org/singleflight"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
 	"github.com/authzed/spicedb/internal/datastore/crdb/migrations"
@@ -278,7 +277,6 @@ type crdbDatastore struct {
 	beginChangefeedQuery string
 	transactionNowQuery  string
 
-	featuresGroup  singleflight.Group[string, *datastore.Features]
 	cachedFeatures *datastore.Features // GUARDED_BY(featuresLock)
 	featuresLock   sync.Mutex
 
@@ -562,9 +560,7 @@ func (cds *crdbDatastore) Features(ctx context.Context) (*datastore.Features, er
 		return cached, nil
 	}
 
-	features, _, err := cds.featuresGroup.Do(ctx, "", func(ictx context.Context) (*datastore.Features, error) {
-		return cds.features(ictx)
-	})
+	features, err := cds.features(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -573,7 +569,7 @@ func (cds *crdbDatastore) Features(ctx context.Context) (*datastore.Features, er
 	cds.cachedFeatures = features
 	cds.featuresLock.Unlock()
 
-	return features, err
+	return features, nil
 }
 
 func (cds *crdbDatastore) features(ctx context.Context) (*datastore.Features, error) {

internal/datastore/proxy/singleflight.go

@@ -2,13 +2,20 @@ package proxy
 
 import (
 	"context"
+	"time"
 
 	"resenje.org/singleflight"
 
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
 )
 
+// singleflightTimeout is the maximum time that a caller of a singleflighted method
+// will wait before giving up on the singleflight executor.
+// This prevents a possible deadlock when all datastore connections are held by goroutines waiting on the
+// singleflight while the singleflight executor is blocked waiting for a connection.
+const singleFlightTimeout = 1 * time.Second
+
 // NewSingleflightDatastoreProxy creates a new Datastore proxy which
 // deduplicates calls to Datastore methods that can share results.
 func NewSingleflightDatastoreProxy(d datastore.Datastore) datastore.Datastore {
@@ -19,6 +26,7 @@ type singleflightProxy struct {
 	headRevGroup  singleflight.Group[string, datastore.RevisionWithSchemaHash]
 	checkRevGroup singleflight.Group[string, string]
 	statsGroup    singleflight.Group[string, datastore.Stats]
+	featuresGroup singleflight.Group[string, *datastore.Features]
 	delegate      datastore.Datastore
 }
 
@@ -43,17 +51,31 @@ func (p *singleflightProxy) ReadWriteTx(ctx context.Context, f datastore.TxUserF
 func (p *singleflightProxy) OptimizedRevision(ctx context.Context) (datastore.RevisionWithSchemaHash, 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, span := tracer.Start(ctx, "singleflightProxy.CheckRevision")
+	defer span.End()
+
+	ctx, cancel := context.WithTimeout(ctx, singleFlightTimeout)
+	defer cancel()
+
 	_, _, err := p.checkRevGroup.Do(ctx, revision.String(), func(ctx context.Context) (string, error) {
 		return "", p.delegate.CheckRevision(ctx, revision)
 	})
 	return err
 }
 
 func (p *singleflightProxy) HeadRevision(ctx context.Context) (datastore.RevisionWithSchemaHash, error) {
+	ctx, span := tracer.Start(ctx, "singleflightProxy.HeadRevision")
+	defer span.End()
+
+	ctx, cancel := context.WithTimeout(ctx, singleFlightTimeout)
+	defer cancel()
+
 	rev, _, err := p.headRevGroup.Do(ctx, "", func(ctx context.Context) (datastore.RevisionWithSchemaHash, error) {
 		return p.delegate.HeadRevision(ctx)
 	})
@@ -69,14 +91,29 @@ func (p *singleflightProxy) Watch(ctx context.Context, afterRevision datastore.R
 }
 
 func (p *singleflightProxy) Statistics(ctx context.Context) (datastore.Stats, error) {
+	ctx, span := tracer.Start(ctx, "singleflightProxy.Statistics")
+	defer span.End()
+
+	ctx, cancel := context.WithTimeout(ctx, singleFlightTimeout)
+	defer cancel()
+
 	stats, _, err := p.statsGroup.Do(ctx, "", func(ctx context.Context) (datastore.Stats, error) {
 		return p.delegate.Statistics(ctx)
 	})
 	return stats, err
 }
 
 func (p *singleflightProxy) Features(ctx context.Context) (*datastore.Features, error) {
-	return p.delegate.Features(ctx)
+	ctx, span := tracer.Start(ctx, "singleflightProxy.Statistics")
+	defer span.End()
+
+	ctx, cancel := context.WithTimeout(ctx, singleFlightTimeout)
+	defer cancel()
+
+	features, _, err := p.featuresGroup.Do(ctx, "", func(ctx context.Context) (*datastore.Features, error) {
+		return p.delegate.Features(ctx)
+	})
+	return features, err
 }
 
 func (p *singleflightProxy) OfflineFeatures() (*datastore.Features, error) {

internal/datastore/revisions/optimized.go

@@ -17,6 +17,12 @@ import (
 	"github.com/authzed/spicedb/pkg/datastore"
 )
 
+// singleflightTimeout is the maximum time that a caller of a singleflighted method
+// will wait before giving up on the singleflight executor.
+// This prevents a possible deadlock when all datastore connections are held by goroutines waiting on the
+// singleflight while the singleflight executor is blocked waiting for a connection.
+const singleFlightTimeout = 1 * time.Second
+
 var tracer = otel.Tracer("spicedb/internal/datastore/common/revisions")
 
 // OptimizedRevisionFunction instructs the datastore to compute its own current
@@ -64,6 +70,9 @@ func (cor *CachedOptimizedRevisions) OptimizedRevision(ctx context.Context) (dat
 	}
 	cor.RUnlock()
 
+	ctx, cancel := context.WithTimeout(ctx, singleFlightTimeout)
+	defer cancel()
+
 	result, _, err := cor.updateGroup.Do(ctx, "", func(ctx context.Context) (datastore.RevisionWithSchemaHash, error) {
 		log.Ctx(ctx).Debug().Time("now", localNow).Msg("computing new revision")
 

pkg/cmd/server/server_test.go

@@ -471,7 +471,7 @@ func TestServerGracefulTerminationOnError(t *testing.T) {
 	}, WithPresharedSecureKey("psk"), WithDatastore(ds), WithEnableMemoryProtectionMiddleware(false))
 	cancel()
 	_, err = c.Complete(ctx)
-	require.NoError(t, err)
+	require.ErrorIs(t, err, context.Canceled)
 }
 
 func TestReplaceUnaryMiddleware(t *testing.T) {

pkg/datalayer/hashcache.go

@@ -16,10 +16,10 @@ import (
 
 var tracer = otel.Tracer("spicedb/pkg/datalayer")
 
-// singleflightTimeout is the maximum time to wait for a singleflight peer to
-// load a schema before falling back to a direct load. This prevents a possible
-// deadlock when all connections in a pool are held by goroutines waiting on the
-// singleflight while the singleflight leader is blocked waiting for a connection.
+// singleflightTimeout is the maximum time that a caller of a singleflighted method
+// will wait before giving up on the singleflight executor.
+// This prevents a possible deadlock when all datastore connections are held by goroutines waiting on the
+// singleflight while the singleflight executor is blocked waiting for a connection.
 //
 //nolint:revive // var instead of const to allow test overrides
 var singleflightTimeout = 1 * time.Second

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) (RevisionWithSchemaHash, 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) (RevisionWithSchemaHash, error) {
+	ctx, span := tracer.Start(ctx, "ctxProxy.HeadRevision")
+	defer span.End()
 	return p.delegate.HeadRevision(context.WithoutCancel(ctx))
 }