scheduler: make ApplyFeatureGates() stateless

pkg/scheduler/algorithmprovider/defaults/defaults.go

@@ -57,7 +57,11 @@ func defaultPredicates() sets.String {
 }
 
 // ApplyFeatureGates applies algorithm by feature gates.
-func ApplyFeatureGates() {
+// The returned function is used to restore the state of registered predicates/priorities
+// when this function is called, and should be called in tests which may modify the value
+// of a feature gate temporarily.
+func ApplyFeatureGates() (restore func()) {
+	snapshot := factory.Copy()
 	if utilfeature.DefaultFeatureGate.Enabled(features.TaintNodesByCondition) {
 		// Remove "CheckNodeCondition", "CheckNodeMemoryPressure", "CheckNodePIDPressure"
 		// and "CheckNodeDiskPressure" predicates
@@ -105,6 +109,11 @@ func ApplyFeatureGates() {
 		// Register the priority function to specific provider too.
 		factory.InsertPriorityKeyToAlgorithmProviderMap(factory.RegisterPriorityMapReduceFunction(priorities.ResourceLimitsPriority, priorities.ResourceLimitsPriorityMap, nil, 1))
 	}
+
+	restore = func() {
+		factory.Apply(snapshot)
+	}
+	return
 }
 
 func registerAlgorithmProvider(predSet, priSet sets.String) {

pkg/scheduler/algorithmprovider/plugins.go

@@ -21,6 +21,6 @@ import (
 )
 
 // ApplyFeatureGates applies algorithm by feature gates.
-func ApplyFeatureGates() {
-	defaults.ApplyFeatureGates()
+func ApplyFeatureGates() func() {
+	return defaults.ApplyFeatureGates()
 }

pkg/scheduler/algorithmprovider/plugins_test.go

@@ -94,7 +94,7 @@ func TestApplyFeatureGates(t *testing.T) {
 	// Apply features for algorithm providers.
 	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.TaintNodesByCondition, true)()
 
-	ApplyFeatureGates()
+	defer ApplyFeatureGates()()
 
 	for _, pn := range algorithmProviderNames {
 		t.Run(pn, func(t *testing.T) {

pkg/scheduler/factory/plugins.go

@@ -101,6 +101,60 @@ type AlgorithmProviderConfig struct {
 	PriorityFunctionKeys sets.String
 }
 
+// Snapshot is used to store current state of registered predicates and priorities.
+type Snapshot struct {
+	fitPredicateMap        map[string]FitPredicateFactory
+	mandatoryFitPredicates sets.String
+	priorityFunctionMap    map[string]PriorityConfigFactory
+	algorithmProviderMap   map[string]AlgorithmProviderConfig
+}
+
+// Copy returns a snapshot of current registered predicates and priorities.
+func Copy() *Snapshot {
+	schedulerFactoryMutex.RLock()
+	defer schedulerFactoryMutex.RUnlock()
+
+	copy := Snapshot{
+		fitPredicateMap:        make(map[string]FitPredicateFactory),
+		mandatoryFitPredicates: sets.NewString(),
+		priorityFunctionMap:    make(map[string]PriorityConfigFactory),
+		algorithmProviderMap:   make(map[string]AlgorithmProviderConfig),
+	}
+	for k, v := range fitPredicateMap {
+		copy.fitPredicateMap[k] = v
+	}
+	for k := range mandatoryFitPredicates {
+		copy.mandatoryFitPredicates[k] = struct{}{}
+	}
+	for k, v := range priorityFunctionMap {
+		copy.priorityFunctionMap[k] = v
+	}
+	for provider, config := range algorithmProviderMap {
+		copyPredKeys, copyPrioKeys := sets.NewString(), sets.NewString()
+		for k := range config.FitPredicateKeys {
+			copyPredKeys[k] = struct{}{}
+		}
+		for k := range config.PriorityFunctionKeys {
+			copyPrioKeys[k] = struct{}{}
+		}
+		copy.algorithmProviderMap[provider] = AlgorithmProviderConfig{
+			FitPredicateKeys:     copyPredKeys,
+			PriorityFunctionKeys: copyPrioKeys,
+		}
+	}
+	return &copy
+}
+
+// Apply sets state of predicates and priorities to `s`.
+func Apply(s *Snapshot) {
+	schedulerFactoryMutex.Lock()
+	fitPredicateMap = s.fitPredicateMap
+	mandatoryFitPredicates = s.mandatoryFitPredicates
+	priorityFunctionMap = s.priorityFunctionMap
+	algorithmProviderMap = s.algorithmProviderMap
+	schedulerFactoryMutex.Unlock()
+}
+
 // RegisterFitPredicate registers a fit predicate with the algorithm
 // registry. Returns the name with which the predicate was registered.
 func RegisterFitPredicate(name string, predicate predicates.FitPredicate) string {

test/integration/daemonset/daemonset_test.go

@@ -89,14 +89,15 @@ func setupScheduler(
 	cs clientset.Interface,
 	informerFactory informers.SharedInformerFactory,
 	stopCh chan struct{},
-) {
+) (restoreFeatureGates func()) {
+	restoreFeatureGates = func() {}
 	// If ScheduleDaemonSetPods is disabled, do not start scheduler.
 	if !utilfeature.DefaultFeatureGate.Enabled(features.ScheduleDaemonSetPods) {
 		return
 	}
 
 	// Enable Features.
-	algorithmprovider.ApplyFeatureGates()
+	restoreFeatureGates = algorithmprovider.ApplyFeatureGates()
 
 	eventBroadcaster := events.NewBroadcaster(&events.EventSinkImpl{
 		Interface: cs.EventsV1beta1().Events(""),
@@ -136,6 +137,7 @@ func setupScheduler(
 	eventBroadcaster.StartRecordingToSink(stopCh)
 
 	go sched.Run()
+	return
 }
 
 func testLabels() map[string]string {
@@ -513,7 +515,7 @@ func TestOneNodeDaemonLaunchesPod(t *testing.T) {
 			defer close(stopCh)
 
 			// Start Scheduler
-			setupScheduler(t, clientset, informers, stopCh)
+			defer setupScheduler(t, clientset, informers, stopCh)()
 
 			informers.Start(stopCh)
 			go dc.Run(5, stopCh)
@@ -557,7 +559,7 @@ func TestSimpleDaemonSetLaunchesPods(t *testing.T) {
 			go dc.Run(5, stopCh)
 
 			// Start Scheduler
-			setupScheduler(t, clientset, informers, stopCh)
+			defer setupScheduler(t, clientset, informers, stopCh)()
 
 			ds := newDaemonSet("foo", ns.Name)
 			ds.Spec.UpdateStrategy = *strategy
@@ -595,7 +597,7 @@ func TestDaemonSetWithNodeSelectorLaunchesPods(t *testing.T) {
 			go dc.Run(5, stopCh)
 
 			// Start Scheduler
-			setupScheduler(t, clientset, informers, stopCh)
+			defer setupScheduler(t, clientset, informers, stopCh)()
 
 			ds := newDaemonSet("foo", ns.Name)
 			ds.Spec.UpdateStrategy = *strategy
@@ -665,7 +667,7 @@ func TestNotReadyNodeDaemonDoesLaunchPod(t *testing.T) {
 		go dc.Run(5, stopCh)
 
 		// Start Scheduler
-		setupScheduler(t, clientset, informers, stopCh)
+		defer setupScheduler(t, clientset, informers, stopCh)()
 
 		ds := newDaemonSet("foo", ns.Name)
 		ds.Spec.UpdateStrategy = *strategy
@@ -753,7 +755,7 @@ func TestInsufficientCapacityNodeWhenScheduleDaemonSetPodsEnabled(t *testing.T)
 		go dc.Run(5, stopCh)
 
 		// Start Scheduler
-		setupScheduler(t, clientset, informers, stopCh)
+		defer setupScheduler(t, clientset, informers, stopCh)()
 
 		ds := newDaemonSet("foo", ns.Name)
 		ds.Spec.Template.Spec = resourcePodSpec("", "120M", "75m")
@@ -816,7 +818,7 @@ func TestLaunchWithHashCollision(t *testing.T) {
 	informers.Start(stopCh)
 	go dc.Run(1, stopCh)
 
-	setupScheduler(t, clientset, informers, stopCh)
+	defer setupScheduler(t, clientset, informers, stopCh)()
 
 	// Create single node
 	_, err := nodeClient.Create(newNode("single-node", nil))
@@ -924,7 +926,7 @@ func TestTaintedNode(t *testing.T) {
 			defer close(stopCh)
 
 			// Start Scheduler
-			setupScheduler(t, clientset, informers, stopCh)
+			defer setupScheduler(t, clientset, informers, stopCh)()
 			informers.Start(stopCh)
 
 			go dc.Run(5, stopCh)
@@ -996,7 +998,7 @@ func TestUnschedulableNodeDaemonDoesLaunchPod(t *testing.T) {
 			go dc.Run(5, stopCh)
 
 			// Start Scheduler
-			setupScheduler(t, clientset, informers, stopCh)
+			defer setupScheduler(t, clientset, informers, stopCh)()
 
 			ds := newDaemonSet("foo", ns.Name)
 			ds.Spec.UpdateStrategy = *strategy

test/integration/scheduler/taint_test.go

@@ -82,7 +82,7 @@ func TestTaintNodeByCondition(t *testing.T) {
 	admission.SetExternalKubeInformerFactory(externalInformers)
 
 	// Apply feature gates to enable TaintNodesByCondition
-	algorithmprovider.ApplyFeatureGates()
+	defer algorithmprovider.ApplyFeatureGates()()
 
 	context = initTestScheduler(t, context, false, nil)
 	cs := context.clientSet