Improve task scheduler rate limiter (#3860)

Author: yycptt

common/dynamicconfig/constants.go

@@ -393,6 +393,11 @@ const (
 
 	// TaskSchedulerEnableRateLimiter indicates if rate limiter should be enabled in task scheduler
 	TaskSchedulerEnableRateLimiter = "history.taskSchedulerEnableRateLimiter"
+	// TaskSchedulerEnableRateLimiterShadowMode indicates if task scheduler rate limiter should run in shadow mode
+	// i.e. through rate limiter and emit metrics but do not actually block/throttle task scheduling
+	TaskSchedulerEnableRateLimiterShadowMode = "history.taskSchedulerEnableRateLimiterShadowMode"
+	// TaskSchedulerThrottleDuration is the throttle duration when task scheduled exceeds max qps
+	TaskSchedulerThrottleDuration = "history.taskSchedulerThrottleDuration"
 	// TaskSchedulerMaxQPS is the max qps task schedulers on a host can schedule tasks
 	// If value less or equal to 0, will fall back to HistoryPersistenceMaxQPS
 	TaskSchedulerMaxQPS = "history.taskSchedulerMaxQPS"

common/metrics/metric_defs.go

@@ -1445,6 +1445,7 @@ var (
 	TaskBatchCompleteCounter                          = NewCounterDef("task_batch_complete_counter")
 	TaskReschedulerPendingTasks                       = NewDimensionlessHistogramDef("task_rescheduler_pending_tasks")
 	PendingTasksCounter                               = NewDimensionlessHistogramDef("pending_tasks")
+	TaskSchedulerThrottled                            = NewCounterDef("task_scheduler_throttled")
 	QueueScheduleLatency                              = NewTimerDef("queue_latency_schedule") // latency for scheduling 100 tasks in one task channel
 	QueueReaderCountHistogram                         = NewDimensionlessHistogramDef("queue_reader_count")
 	QueueSliceCountHistogram                          = NewDimensionlessHistogramDef("queue_slice_count")

common/tasks/benchmark_test.go

@@ -33,6 +33,7 @@ import (
 	"go.temporal.io/server/common/clock"
 	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/log"
+	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/quotas"
 )
 
@@ -62,15 +63,18 @@ func BenchmarkInterleavedWeightedRoundRobinScheduler_Sequential(b *testing.B) {
 
 	scheduler := NewInterleavedWeightedRoundRobinScheduler(
 		InterleavedWeightedRoundRobinSchedulerOptions[*noopTask, int]{
-			TaskChannelKeyFn:      func(nt *noopTask) int { return rand.Intn(4) },
-			ChannelWeightFn:       func(key int) int { return channelKeyToWeight[key] },
-			ChannelQuotaRequestFn: func(key int) quotas.Request { return quotas.NewRequest("", 1, "", "", "") },
-			EnableRateLimiter:     dynamicconfig.GetBoolPropertyFn(true),
+			TaskChannelKeyFn:            func(nt *noopTask) int { return rand.Intn(4) },
+			ChannelWeightFn:             func(key int) int { return channelKeyToWeight[key] },
+			ChannelQuotaRequestFn:       func(key int) quotas.Request { return quotas.NewRequest("", 1, "", "", "") },
+			TaskChannelMetricTagsFn:     func(key int) []metrics.Tag { return nil },
+			EnableRateLimiter:           dynamicconfig.GetBoolPropertyFn(true),
+			EnableRateLimiterShadowMode: dynamicconfig.GetBoolPropertyFn(false),
 		},
 		Scheduler[*noopTask](&noopScheduler{}),
 		quotas.NoopRequestRateLimiter,
 		clock.NewRealTimeSource(),
 		logger,
+		metrics.NoopMetricsHandler,
 	)
 	scheduler.Start()
 	defer scheduler.Stop()
@@ -92,15 +96,18 @@ func BenchmarkInterleavedWeightedRoundRobinScheduler_Parallel(b *testing.B) {
 
 	scheduler := NewInterleavedWeightedRoundRobinScheduler(
 		InterleavedWeightedRoundRobinSchedulerOptions[*noopTask, int]{
-			TaskChannelKeyFn:      func(nt *noopTask) int { return rand.Intn(4) },
-			ChannelWeightFn:       func(key int) int { return channelKeyToWeight[key] },
-			ChannelQuotaRequestFn: func(key int) quotas.Request { return quotas.NewRequest("", 1, "", "", "") },
-			EnableRateLimiter:     dynamicconfig.GetBoolPropertyFn(true),
+			TaskChannelKeyFn:            func(nt *noopTask) int { return rand.Intn(4) },
+			ChannelWeightFn:             func(key int) int { return channelKeyToWeight[key] },
+			ChannelQuotaRequestFn:       func(key int) quotas.Request { return quotas.NewRequest("", 1, "", "", "") },
+			TaskChannelMetricTagsFn:     func(key int) []metrics.Tag { return nil },
+			EnableRateLimiter:           dynamicconfig.GetBoolPropertyFn(true),
+			EnableRateLimiterShadowMode: dynamicconfig.GetBoolPropertyFn(false),
 		},
 		Scheduler[*noopTask](&noopScheduler{}),
 		quotas.NoopRequestRateLimiter,
 		clock.NewRealTimeSource(),
 		logger,
+		metrics.NoopMetricsHandler,
 	)
 	scheduler.Start()
 	defer scheduler.Stop()