feat: Resource rate limit (#1433)

BEGIN_COMMIT_OVERRIDE
feat: Introduce a per resource rate limit in addition to a global resource rate limit

feat: Set default rate limit of `5` for `SingleResourceMaxConcurrency` and `SingleNestedTableMaxConcurrency`
END_COMMIT_OVERRIDE

This introduces a rate limit on a resource + table + client basis. Prior to this we only had rate limits on a per table+ client basis and all resources shared the same concurrency bucket

This also adds a default value of 5 for both `singleResourceMaxConcurrency` and  `singleNestedTableMaxConcurrency`

Author: bbernays

scheduler/scheduler.go

@@ -25,11 +25,13 @@ import (
 )
 
 const (
-	DefaultConcurrency     = 50000
-	DefaultMaxDepth        = 4
-	minTableConcurrency    = 1
-	minResourceConcurrency = 100
-	otelName               = "schedule"
+	DefaultSingleResourceMaxConcurrency    = 5
+	DefaultSingleNestedTableMaxConcurrency = 5
+	DefaultConcurrency                     = 50000
+	DefaultMaxDepth                        = 4
+	minTableConcurrency                    = 1
+	minResourceConcurrency                 = 100
+	otelName                               = "schedule"
 )
 
 var ErrNoTables = errors.New("no tables specified for syncing, review `tables` and `skip_tables` in your config and specify at least one table to sync")
@@ -72,6 +74,12 @@ func WithSingleNestedTableMaxConcurrency(concurrency int64) Option {
 	}
 }
 
+func WithSingleResourceMaxConcurrency(concurrency int64) Option {
+	return func(s *Scheduler) {
+		s.singleResourceMaxConcurrency = concurrency
+	}
+}
+
 type SyncOption func(*syncClient)
 
 func WithSyncDeterministicCQID(deterministicCQID bool) SyncOption {
@@ -99,6 +107,9 @@ type Scheduler struct {
 	singleTableConcurrency sync.Map
 	// The maximum number of go routines that can be spawned for a single table+client pair
 	singleNestedTableMaxConcurrency int64
+
+	// The maximum number of go routines that can be spawned for a specific resource
+	singleResourceMaxConcurrency int64
 }
 
 type syncClient struct {
@@ -113,9 +124,11 @@ type syncClient struct {
 
 func NewScheduler(opts ...Option) *Scheduler {
 	s := Scheduler{
-		caser:       caser.New(),
-		concurrency: DefaultConcurrency,
-		maxDepth:    DefaultMaxDepth,
+		caser:                           caser.New(),
+		concurrency:                     DefaultConcurrency,
+		maxDepth:                        DefaultMaxDepth,
+		singleResourceMaxConcurrency:    DefaultSingleResourceMaxConcurrency,
+		singleNestedTableMaxConcurrency: DefaultSingleNestedTableMaxConcurrency,
 	}
 	for _, opt := range opts {
 		opt(&s)
@@ -132,10 +145,6 @@ func NewScheduler(opts ...Option) *Scheduler {
 	}
 	s.resourceSem = semaphore.NewWeighted(int64(resourceConcurrency))
 
-	// To preserve backwards compatibility, if singleTableMaxConcurrency is not set, set it to the max concurrency
-	if s.singleNestedTableMaxConcurrency == 0 {
-		s.singleNestedTableMaxConcurrency = int64(tableConcurrency)
-	}
 	return &s
 }
 

scheduler/scheduler_dfs.go

@@ -135,14 +135,28 @@ func (s *syncClient) resolveResourcesDfs(ctx context.Context, table *schema.Tabl
 		sentValidationErrors := sync.Map{}
 		for i := range resourcesSlice {
 			i := i
-			if err := s.scheduler.resourceSem.Acquire(ctx, 1); err != nil {
+			resourceConcurrencyKey := table.Name + "-" + client.ID() + "-" + "resource"
+			resourceSemVal, _ := s.scheduler.singleTableConcurrency.LoadOrStore(resourceConcurrencyKey, semaphore.NewWeighted(s.scheduler.singleResourceMaxConcurrency))
+			resourceSem := resourceSemVal.(*semaphore.Weighted)
+			if err := resourceSem.Acquire(ctx, 1); err != nil {
 				s.logger.Warn().Err(err).Msg("failed to acquire semaphore. context cancelled")
+				// This means context was cancelled
+				wg.Wait()
+				// we have to continue emptying the channel to exit gracefully
+				return
+			}
+
+			// Once Resource semaphore is acquired we can acquire the global semaphore
+			if err := s.scheduler.resourceSem.Acquire(ctx, 1); err != nil {
+				// This means context was cancelled
+				resourceSem.Release(1)
 				wg.Wait()
 				// we have to continue emptying the channel to exit gracefully
 				return
 			}
 			wg.Add(1)
 			go func() {
+				defer resourceSem.Release(1)
 				defer s.scheduler.resourceSem.Release(1)
 				defer wg.Done()
 				//nolint:all