add code for auto enabling priority and fairness

Author: majiru

api/persistence/v1/task_queues.pb.go

@@ -1382,6 +1382,11 @@ second per poller by one physical queue manager`,
 The metric has 2 dimensions: namespace_id and plugin_name. Disabled by default as this is 
 an optional feature and also requires a metrics collection system that can handle higher cardinalities.`,
 	)
+	MatchingAutoEnable = NewTaskQueueBoolSetting(
+		"matching.autoEnable",
+		false,
+		`MatchingAutoEnable automatically enables fairness when a fairness key is seen`,
+	)
 
 	// keys for history
 

common/dynamicconfig/constants.go

@@ -134,6 +134,13 @@ message TaskQueueTypeUserData {
     DeploymentData deployment_data = 1;
 
     temporal.api.taskqueue.v1.TaskQueueConfig config = 2;
+
+    enum FairnessState {
+        FAIRNESS_STATE_UNSPECIFIED = 0;
+        FAIRNESS_STATE_V1 = 1;
+        FAIRNESS_STATE_V2 = 2;
+    };
+    FairnessState fairness_state = 3;
 }
 
 // Container for all persistent user provided data for a task queue family.

proto/internal/temporal/server/api/persistence/v1/task_queues.proto

@@ -82,6 +82,7 @@ type (
 		MembershipUnloadDelay                    dynamicconfig.DurationPropertyFn
 		TaskQueueInfoByBuildIdTTL                dynamicconfig.DurationPropertyFnWithTaskQueueFilter
 		PriorityLevels                           dynamicconfig.IntPropertyFnWithTaskQueueFilter
+		AutoEnableSub                            dynamicconfig.TypedSubscribableWithTaskQueueFilter[bool]
 
 		RateLimiterRefreshInterval    time.Duration
 		FairnessKeyRateLimitCacheSize dynamicconfig.IntPropertyFnWithTaskQueueFilter
@@ -149,6 +150,8 @@ type (
 		EnableFairness             bool
 		EnableFairnessSub          func(func(bool)) (bool, func())
 		EnableMigration            func() bool
+		AutoEnable                 bool
+		AutoEnableSub              func(func(bool)) (bool, func())
 		GetTasksBatchSize          func() int
 		GetTasksReloadAt           func() int
 		UpdateAckInterval          func() time.Duration
@@ -308,6 +311,7 @@ func NewConfig(
 		FairnessKeyRateLimitCacheSize:            dynamicconfig.MatchingFairnessKeyRateLimitCacheSize.Get(dc),
 		MaxFairnessKeyWeightOverrides:            dynamicconfig.MatchingMaxFairnessKeyWeightOverrides.Get(dc),
 		MaxIDLengthLimit:                         dynamicconfig.MaxIDLengthLimit.Get(dc),
+		AutoEnableSub:                            dynamicconfig.MatchingAutoEnable.Subscribe(dc),
 
 		AdminNamespaceToPartitionDispatchRate:          dynamicconfig.AdminMatchingNamespaceToPartitionDispatchRate.Get(dc),
 		AdminNamespaceToPartitionRateSub:               dynamicconfig.AdminMatchingNamespaceToPartitionDispatchRate.Subscribe(dc),
@@ -355,6 +359,9 @@ func newTaskQueueConfig(tq *tqid.TaskQueue, config *Config, ns namespace.Name) *
 		EnableMigration: func() bool {
 			return config.EnableMigration(ns.String(), taskQueueName, taskType)
 		},
+		AutoEnableSub: func(cb func(bool)) (bool, func()) {
+			return config.AutoEnableSub(ns.String(), taskQueueName, taskType, cb)
+		},
 		GetTasksBatchSize: func() int {
 			return config.GetTasksBatchSize(ns.String(), taskQueueName, taskType)
 		},

service/matching/config.go

@@ -430,7 +430,11 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 	tqConfig.loadCause = loadCause
 	logger, throttledLogger, metricsHandler := e.loggerAndMetricsForPartition(namespaceEntry, partition, tqConfig)
 	onFatalErr := func(cause unloadCause) { newPM.unloadFromEngine(cause) }
-	onUserDataChanged := func() { newPM.userDataChanged() }
+	done := make(chan struct{})
+	onUserDataChanged := func(old, new *persistencespb.VersionedTaskQueueUserData) {
+		<-done
+		newPM.userDataChanged(old, new)
+	}
 	userDataManager := newUserDataManager(
 		e.taskManager,
 		e.matchingRawClient,
@@ -442,6 +446,7 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 		e.namespaceRegistry,
 	)
 	newPM, err = newTaskQueuePartitionManager(
+		ctx,
 		e,
 		namespaceEntry,
 		partition,
@@ -451,6 +456,7 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 		metricsHandler,
 		userDataManager,
 	)
+	close(done)
 	if err != nil {
 		return nil, false, err
 	}

service/matching/matching_engine.go

@@ -267,8 +267,10 @@ func newMatchingEngine(
 func (s *matchingEngineSuite) newPartitionManager(prtn tqid.Partition, config *Config) taskQueuePartitionManager {
 	tqConfig := newTaskQueueConfig(prtn.TaskQueue(), config, matchingTestNamespace)
 	logger, _, metricsHandler := s.matchingEngine.loggerAndMetricsForPartition(s.ns, prtn, tqConfig)
-	pm, err := newTaskQueuePartitionManager(s.matchingEngine, s.ns, prtn, tqConfig, logger, logger, metricsHandler, &mockUserDataManager{})
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	pm, err := newTaskQueuePartitionManager(ctx, s.matchingEngine, s.ns, prtn, tqConfig, logger, logger, metricsHandler, &mockUserDataManager{})
 	s.Require().NoError(err)
+	cancel()
 	return pm
 }
 

service/matching/matching_engine_test.go

@@ -87,8 +87,10 @@ func (s *PhysicalTaskQueueManagerTestSuite) SetupTest() {
 	onFatalErr := func(unloadCause) { s.T().Fatal("user data manager called onFatalErr") }
 	udMgr := newUserDataManager(engine.taskManager, engine.matchingRawClient, onFatalErr, nil, prtn, tqConfig, engine.logger, engine.namespaceRegistry)
 
-	prtnMgr, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	prtnMgr, err := newTaskQueuePartitionManager(ctx, engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)
 	s.NoError(err)
+	cancel()
 	engine.partitions[prtn.Key()] = prtnMgr
 
 	s.tqMgr, err = newPhysicalTaskQueueManager(prtnMgr, s.physicalTaskQueueKey)

service/matching/physical_task_queue_manager_test.go

@@ -15,6 +15,7 @@ import (
 	"go.temporal.io/server/api/matchingservice/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	taskqueuespb "go.temporal.io/server/api/taskqueue/v1"
+	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/cache"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
@@ -63,8 +64,11 @@ type (
 		// TODO(stephanos): move cache out of partition manager
 		cache cache.Cache // non-nil for root-partition
 
+		fairnessState persistencespb.TaskQueueTypeUserData_FairnessState //Set once on initialization and read only after
+
 		cancelNewMatcherSub func()
 		cancelFairnessSub   func()
+		cancelAutoEnableSub func()
 
 		// rateLimitManager is used to manage the rate limit for task queues.
 		rateLimitManager *rateLimitManager
@@ -82,6 +86,7 @@ func (pm *taskQueuePartitionManagerImpl) GetCache(key any) any {
 var _ taskQueuePartitionManager = (*taskQueuePartitionManagerImpl)(nil)
 
 func newTaskQueuePartitionManager(
+	ctx context.Context,
 	e *matchingEngineImpl,
 	ns *namespace.Namespace,
 	partition tqid.Partition,
@@ -119,14 +124,40 @@ func newTaskQueuePartitionManager(
 		pm.unloadFromEngine(unloadCauseConfigChange)
 	}
 
-	var fairness bool
-	fairness, pm.cancelFairnessSub = tqConfig.EnableFairnessSub(unload)
-	// Fairness is disabled for sticky queues for now so that we can still use TTLs.
-	tqConfig.EnableFairness = fairness && partition.Kind() != enumspb.TASK_QUEUE_KIND_STICKY
-	if fairness {
+	pm.userDataManager.Start()
+	err := pm.userDataManager.WaitUntilInitialized(ctx)
+	if err != nil {
+		return nil, err
+	}
+	//todo(moody): do we need to be more cautious loading this? Do we need to ensure that we have PerType()? Check for a nil return?
+	data, _, err := pm.getPerTypeUserData()
+	if err != nil {
+		return nil, err
+	}
+
+	tqConfig.AutoEnable, pm.cancelAutoEnableSub = tqConfig.AutoEnableSub(unload)
+	pm.fairnessState = data.GetFairnessState()
+	switch {
+	case tqConfig.AutoEnable == false || pm.fairnessState == persistencespb.TaskQueueTypeUserData_FAIRNESS_STATE_UNSPECIFIED:
+		var fairness bool
+		fairness, pm.cancelFairnessSub = tqConfig.EnableFairnessSub(unload)
+		// Fairness is disabled for sticky queues for now so that we can still use TTLs.
+		tqConfig.EnableFairness = fairness && partition.Kind() != enumspb.TASK_QUEUE_KIND_STICKY
+		if fairness {
+			tqConfig.NewMatcher = true
+		} else {
+			tqConfig.NewMatcher, pm.cancelNewMatcherSub = tqConfig.NewMatcherSub(unload)
+		}
+	case pm.fairnessState == persistencespb.TaskQueueTypeUserData_FAIRNESS_STATE_V1:
 		tqConfig.NewMatcher = true
-	} else {
-		tqConfig.NewMatcher, pm.cancelNewMatcherSub = tqConfig.NewMatcherSub(unload)
+		tqConfig.EnableFairness = false
+	case pm.fairnessState == persistencespb.TaskQueueTypeUserData_FAIRNESS_STATE_V2:
+		tqConfig.NewMatcher = true
+		if partition.Kind() == enumspb.TASK_QUEUE_KIND_STICKY {
+			tqConfig.EnableFairness = false
+		} else {
+			tqConfig.EnableFairness = true
+		}
 	}
 
 	defaultQ, err := newPhysicalTaskQueueManager(pm, UnversionedQueueKey(partition))
@@ -139,7 +170,6 @@ func newTaskQueuePartitionManager(
 
 func (pm *taskQueuePartitionManagerImpl) Start() {
 	pm.engine.updateTaskQueuePartitionGauge(pm.Namespace(), pm.partition, 1)
-	pm.userDataManager.Start()
 	pm.defaultQueue.Start()
 }
 
@@ -159,6 +189,9 @@ func (pm *taskQueuePartitionManagerImpl) Stop(unloadCause unloadCause) {
 	if pm.cancelNewMatcherSub != nil {
 		pm.cancelNewMatcherSub()
 	}
+	if pm.cancelAutoEnableSub != nil {
+		pm.cancelAutoEnableSub()
+	}
 
 	// First, stop all queues to wrap up ongoing operations.
 	for _, vq := range pm.versionedQueues {
@@ -184,10 +217,6 @@ func (pm *taskQueuePartitionManagerImpl) MarkAlive() {
 }
 
 func (pm *taskQueuePartitionManagerImpl) WaitUntilInitialized(ctx context.Context) error {
-	err := pm.userDataManager.WaitUntilInitialized(ctx)
-	if err != nil {
-		return err
-	}
 	return pm.defaultQueue.WaitUntilInitialized(ctx)
 }
 
@@ -197,6 +226,21 @@ func (pm *taskQueuePartitionManagerImpl) AddTask(
 ) (buildId string, syncMatched bool, err error) {
 	var spoolQueue, syncMatchQueue physicalTaskQueueManager
 	directive := params.taskInfo.GetVersionDirective()
+
+	if pm.Partition().IsRoot() && pm.config.AutoEnable && pm.fairnessState == persistencespb.TaskQueueTypeUserData_FAIRNESS_STATE_UNSPECIFIED {
+		//TODO(moody): unsure about this check, what is the correct way to check to see if PriorityKey is set?
+		//TODO(moody): This was originally discussed to be a separate API, but is just exposing this through the generic UpdateUserData sufficient?
+		// what is the pro/cons of adding the new API and perhaps invoking that instead? We're going to unload either way...
+		if params.taskInfo.Priority != nil && (params.taskInfo.Priority.FairnessKey != "" || params.taskInfo.Priority.PriorityKey != int32(pm.config.DefaultPriorityKey)) {
+			updateFn := func(old *persistencespb.TaskQueueUserData) (*persistencespb.TaskQueueUserData, bool, error) {
+				new := common.CloneProto(old)
+				perType := new.GetPerType()[int32(pm.Partition().TaskType())]
+				perType.FairnessState = persistencespb.TaskQueueTypeUserData_FAIRNESS_STATE_V2
+				return new, true, nil
+			}
+			pm.userDataManager.UpdateUserData(ctx, UserDataUpdateOptions{Source: "Matching auto enable"}, updateFn)
+		}
+	}
 	// spoolQueue will be nil iff task is forwarded.
 reredirectTask:
 	spoolQueue, syncMatchQueue, _, taskDispatchRevisionNumber, err := pm.getPhysicalQueuesForAdd(ctx, directive, params.forwardInfo, params.taskInfo.GetRunId(), params.taskInfo.GetWorkflowId(), false)
@@ -1303,7 +1347,17 @@ func (pm *taskQueuePartitionManagerImpl) getPerTypeUserData() (*persistencespb.T
 	return perType, userDataChanged, nil
 }
 
-func (pm *taskQueuePartitionManagerImpl) userDataChanged() {
+func (pm *taskQueuePartitionManagerImpl) userDataChanged(old, new *persistencespb.VersionedTaskQueueUserData) {
+	taskType := int32(pm.Partition().TaskType())
+	//TODO(moody): this stinks, do we need this more verbose, is this interface bad?
+	//TODO(moody): we get calls into this callback quite a bit with data being nil(sampled from unit tests), do we want to go through the full update
+	// even in those cases? Should we pass the inner data and avoid a callback when that is nil? I am not totally sure about the implcations....
+	if old != nil && old.GetData() != nil && old.GetData().GetPerType() != nil && new != nil && new.GetData() != nil && new.GetData().GetPerType() != nil {
+		if old.GetData().GetPerType()[taskType].FairnessState != new.GetData().GetPerType()[taskType].FairnessState {
+			pm.unloadFromEngine(unloadCauseConfigChange)
+			return
+		}
+	}
 	// Update rateLimits if any change is userData.
 	pm.rateLimitManager.UserDataChanged()
 

service/matching/task_queue_partition_manager.go

@@ -85,7 +85,7 @@ func (s *PartitionManagerTestSuite) SetupTest() {
 	tqConfig := newTaskQueueConfig(partition.TaskQueue(), engine.config, ns.Name())
 	s.userDataMgr = &mockUserDataManager{}
 
-	pm, err := newTaskQueuePartitionManager(engine, ns, partition, tqConfig, logger, logger, metrics.NoopMetricsHandler, s.userDataMgr)
+	pm, err := newTaskQueuePartitionManager(context.Background(), engine, ns, partition, tqConfig, logger, logger, metrics.NoopMetricsHandler, s.userDataMgr)
 	s.NoError(err)
 	s.partitionMgr = pm
 	engine.Start()

service/matching/task_queue_partition_manager_test.go

@@ -75,7 +75,7 @@ type (
 	userDataManagerImpl struct {
 		lock              sync.Mutex
 		onFatalErr        func(unloadCause)
-		onUserDataChanged func() // if set, call this in new goroutine when user data changes
+		onUserDataChanged func(old, new *persistencespb.VersionedTaskQueueUserData) // if set, call this in new goroutine when user data changes
 		partition         tqid.Partition
 		userData          *persistencespb.VersionedTaskQueueUserData
 		userDataChanged   chan struct{}
@@ -110,7 +110,7 @@ func newUserDataManager(
 	store persistence.TaskManager,
 	matchingClient matchingservice.MatchingServiceClient,
 	onFatalErr func(unloadCause),
-	onUserDataChanged func(),
+	onUserDataChanged func(old, new *persistencespb.VersionedTaskQueueUserData),
 	partition tqid.Partition,
 	config *taskQueueConfig,
 	logger log.Logger,
@@ -176,11 +176,12 @@ func (m *userDataManagerImpl) getUserDataLocked() (*persistencespb.VersionedTask
 }
 
 func (m *userDataManagerImpl) setUserDataLocked(userData *persistencespb.VersionedTaskQueueUserData) {
+	old := m.userData
 	m.userData = userData
 	close(m.userDataChanged)
 	m.userDataChanged = make(chan struct{})
 	if m.onUserDataChanged != nil {
-		go m.onUserDataChanged()
+		go m.onUserDataChanged(old, userData)
 	}
 }