Merge pull request #1997 from authzed/dispatch-chunk-safeguards

additional dispatch chunk safeguards

Author: vroldanbet

internal/datastore/common/sql.go

@@ -11,6 +11,7 @@ import (
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 
+	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -107,6 +108,11 @@ type SchemaQueryFilterer struct {
 
 // NewSchemaQueryFilterer creates a new SchemaQueryFilterer object.
 func NewSchemaQueryFilterer(schema SchemaInformation, initialQuery sq.SelectBuilder, filterMaximumIDCount uint16) SchemaQueryFilterer {
+	if filterMaximumIDCount == 0 {
+		filterMaximumIDCount = 100
+		log.Warn().Msg("SchemaQueryFilterer: filterMaximumIDCount not set, defaulting to 100")
+	}
+
 	return SchemaQueryFilterer{
 		schema:                schema,
 		queryBuilder:          initialQuery,

internal/datastore/crdb/options.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
+	log "github.com/authzed/spicedb/internal/logging"
 )
 
 type crdbOptions struct {
@@ -85,6 +86,11 @@ func generateConfig(options []Option) (crdbOptions, error) {
 		)
 	}
 
+	if computed.filterMaximumIDCount == 0 {
+		computed.filterMaximumIDCount = 100
+		log.Warn().Msg("filterMaximumIDCount not set, defaulting to 100")
+	}
+
 	return computed, nil
 }
 

internal/datastore/mysql/options.go

@@ -3,6 +3,8 @@ package mysql
 import (
 	"fmt"
 	"time"
+
+	log "github.com/authzed/spicedb/internal/logging"
 )
 
 const (
@@ -82,6 +84,11 @@ func generateConfig(options []Option) (mysqlOptions, error) {
 		)
 	}
 
+	if computed.filterMaximumIDCount == 0 {
+		computed.filterMaximumIDCount = 100
+		log.Warn().Msg("filterMaximumIDCount not set, defaulting to 100")
+	}
+
 	return computed, nil
 }
 

internal/datastore/postgres/options.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
+	log "github.com/authzed/spicedb/internal/logging"
 )
 
 type postgresOptions struct {
@@ -106,6 +107,11 @@ func generateConfig(options []Option) (postgresOptions, error) {
 		return computed, fmt.Errorf("unknown migration phase: %s", computed.migrationPhase)
 	}
 
+	if computed.filterMaximumIDCount == 0 {
+		computed.filterMaximumIDCount = 100
+		log.Warn().Msg("filterMaximumIDCount not set, defaulting to 100")
+	}
+
 	return computed, nil
 }
 

internal/datastore/spanner/options.go

@@ -5,6 +5,8 @@ import (
 	"math"
 	"runtime"
 	"time"
+
+	log "github.com/authzed/spicedb/internal/logging"
 )
 
 type spannerOptions struct {
@@ -88,6 +90,11 @@ func generateConfig(options []Option) (spannerOptions, error) {
 		return computed, fmt.Errorf("unknown migration phase: %s", computed.migrationPhase)
 	}
 
+	if computed.filterMaximumIDCount == 0 {
+		computed.filterMaximumIDCount = 100
+		log.Warn().Msg("filterMaximumIDCount not set, defaulting to 100")
+	}
+
 	return computed, nil
 }
 

internal/dispatch/cluster/cluster.go

@@ -7,6 +7,7 @@ import (
 	"github.com/authzed/spicedb/internal/dispatch/caching"
 	"github.com/authzed/spicedb/internal/dispatch/graph"
 	"github.com/authzed/spicedb/internal/dispatch/keys"
+	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/cache"
 )
 
@@ -74,6 +75,11 @@ func NewClusterDispatcher(dispatch dispatch.Dispatcher, options ...Option) (disp
 		fn(&opts)
 	}
 
+	chunkSize := opts.dispatchChunkSize
+	if chunkSize == 0 {
+		chunkSize = 100
+		log.Warn().Msgf("ClusterDispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
+	}
 	clusterDispatch := graph.NewDispatcher(dispatch, opts.concurrencyLimits, opts.dispatchChunkSize)
 
 	if opts.prometheusSubsystem == "" {

internal/dispatch/combined/combined.go

@@ -153,6 +153,7 @@ func NewDispatcher(options ...Option) (dispatch.Dispatcher, error) {
 	chunkSize := opts.dispatchChunkSize
 	if chunkSize == 0 {
 		chunkSize = 100
+		log.Warn().Msgf("CombinedDispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
 	}
 	redispatch := graph.NewDispatcher(cachingRedispatch, opts.concurrencyLimits, chunkSize)
 	redispatch = singleflight.New(redispatch, &keys.CanonicalKeyHandler{})

internal/dispatch/graph/graph.go

@@ -88,13 +88,18 @@ func NewLocalOnlyDispatcherWithLimits(concurrencyLimits ConcurrencyLimits, dispa
 	d := &localDispatcher{}
 
 	concurrencyLimits = limitsOrDefaults(concurrencyLimits, defaultConcurrencyLimit)
+	chunkSize := dispatchChunkSize
+	if chunkSize == 0 {
+		chunkSize = 100
+		log.Warn().Msgf("LocalOnlyDispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
+	}
 
-	d.checker = graph.NewConcurrentChecker(d, concurrencyLimits.Check, dispatchChunkSize)
+	d.checker = graph.NewConcurrentChecker(d, concurrencyLimits.Check, chunkSize)
 	d.expander = graph.NewConcurrentExpander(d)
-	d.reachableResourcesHandler = graph.NewCursoredReachableResources(d, concurrencyLimits.ReachableResources, dispatchChunkSize)
-	d.lookupResourcesHandler = graph.NewCursoredLookupResources(d, d, concurrencyLimits.LookupResources, dispatchChunkSize)
-	d.lookupSubjectsHandler = graph.NewConcurrentLookupSubjects(d, concurrencyLimits.LookupSubjects, dispatchChunkSize)
-	d.lookupResourcesHandler2 = graph.NewCursoredLookupResources2(d, d, concurrencyLimits.LookupResources, dispatchChunkSize)
+	d.reachableResourcesHandler = graph.NewCursoredReachableResources(d, concurrencyLimits.ReachableResources, chunkSize)
+	d.lookupResourcesHandler = graph.NewCursoredLookupResources(d, d, concurrencyLimits.LookupResources, chunkSize)
+	d.lookupSubjectsHandler = graph.NewConcurrentLookupSubjects(d, concurrencyLimits.LookupSubjects, chunkSize)
+	d.lookupResourcesHandler2 = graph.NewCursoredLookupResources2(d, d, concurrencyLimits.LookupResources, chunkSize)
 
 	return d
 }
@@ -103,13 +108,18 @@ func NewLocalOnlyDispatcherWithLimits(concurrencyLimits ConcurrencyLimits, dispa
 // the provided redispatcher.
 func NewDispatcher(redispatcher dispatch.Dispatcher, concurrencyLimits ConcurrencyLimits, dispatchChunkSize uint16) dispatch.Dispatcher {
 	concurrencyLimits = limitsOrDefaults(concurrencyLimits, defaultConcurrencyLimit)
+	chunkSize := dispatchChunkSize
+	if chunkSize == 0 {
+		chunkSize = 100
+		log.Warn().Msgf("Dispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
+	}
 
-	checker := graph.NewConcurrentChecker(redispatcher, concurrencyLimits.Check, dispatchChunkSize)
+	checker := graph.NewConcurrentChecker(redispatcher, concurrencyLimits.Check, chunkSize)
 	expander := graph.NewConcurrentExpander(redispatcher)
-	reachableResourcesHandler := graph.NewCursoredReachableResources(redispatcher, concurrencyLimits.ReachableResources, dispatchChunkSize)
-	lookupResourcesHandler := graph.NewCursoredLookupResources(redispatcher, redispatcher, concurrencyLimits.LookupResources, dispatchChunkSize)
-	lookupSubjectsHandler := graph.NewConcurrentLookupSubjects(redispatcher, concurrencyLimits.LookupSubjects, dispatchChunkSize)
-	lookupResourcesHandler2 := graph.NewCursoredLookupResources2(redispatcher, redispatcher, concurrencyLimits.LookupResources, dispatchChunkSize)
+	reachableResourcesHandler := graph.NewCursoredReachableResources(redispatcher, concurrencyLimits.ReachableResources, chunkSize)
+	lookupResourcesHandler := graph.NewCursoredLookupResources(redispatcher, redispatcher, concurrencyLimits.LookupResources, chunkSize)
+	lookupSubjectsHandler := graph.NewConcurrentLookupSubjects(redispatcher, concurrencyLimits.LookupSubjects, chunkSize)
+	lookupResourcesHandler2 := graph.NewCursoredLookupResources2(redispatcher, redispatcher, concurrencyLimits.LookupResources, chunkSize)
 
 	return &localDispatcher{
 		checker:                   checker,

internal/services/v1/experimental.go

@@ -77,6 +77,14 @@ func NewExperimentalServer(dispatch dispatch.Dispatcher, permServerConfig Permis
 		config.StreamReadTimeout = streamReadTimeoutFallbackSeconds * time.Second
 	}
 
+	chunkSize := permServerConfig.DispatchChunkSize
+	if chunkSize == 0 {
+		log.
+			Warn().
+			Msg("experimental server config specified invalid DispatchChunkSize, defaulting to 100")
+		chunkSize = 100
+	}
+
 	return &experimentalServer{
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
@@ -98,7 +106,7 @@ func NewExperimentalServer(dispatch dispatch.Dispatcher, permServerConfig Permis
 			maxCaveatContextSize: permServerConfig.MaxCaveatContextSize,
 			maxConcurrency:       config.BulkCheckMaxConcurrency,
 			dispatch:             dispatch,
-			dispatchChunkSize:    permServerConfig.DispatchChunkSize,
+			dispatchChunkSize:    chunkSize,
 		},
 	}
 }

pkg/genutil/slicez/chunking.go

@@ -14,8 +14,8 @@ func ForEachChunk[T any](data []T, chunkSize uint16, handler func(items []T)) {
 
 func ForEachChunkUntil[T any](data []T, chunkSize uint16, handler func(items []T) (bool, error)) (bool, error) {
 	if chunkSize == 0 {
-		logging.Warn().Int("invalid-chunk-size", int(chunkSize)).Msg("ForEachChunk got an invalid chunk size; defaulting to 1")
-		chunkSize = 1
+		logging.Warn().Int("invalid-chunk-size", int(chunkSize)).Msg("ForEachChunk got an invalid chunk size; defaulting to 100")
+		chunkSize = 100
 	}
 
 	dataLength := uint64(len(data))