diff --git a/pkg/scheduler/algorithmprovider/defaults/defaults.go b/pkg/scheduler/algorithmprovider/defaults/defaults.go index 6fd14d4513f..a37ec478e87 100644 --- a/pkg/scheduler/algorithmprovider/defaults/defaults.go +++ b/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) { diff --git a/pkg/scheduler/algorithmprovider/plugins.go b/pkg/scheduler/algorithmprovider/plugins.go index e2784f62609..8af8f6fa099 100644 --- a/pkg/scheduler/algorithmprovider/plugins.go +++ b/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() } diff --git a/pkg/scheduler/algorithmprovider/plugins_test.go b/pkg/scheduler/algorithmprovider/plugins_test.go index 811e4ca4edf..4ae08651083 100644 --- a/pkg/scheduler/algorithmprovider/plugins_test.go +++ b/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) { diff --git a/pkg/scheduler/factory/plugins.go b/pkg/scheduler/factory/plugins.go index 1bf8fe81e3a..ecbc772d7d1 100644 --- a/pkg/scheduler/factory/plugins.go +++ b/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 © +} + +// 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 { diff --git a/test/integration/daemonset/daemonset_test.go b/test/integration/daemonset/daemonset_test.go index f61a2dc66f3..e62a995fc2d 100644 --- a/test/integration/daemonset/daemonset_test.go +++ b/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 diff --git a/test/integration/scheduler/taint_test.go b/test/integration/scheduler/taint_test.go index 422a1cd2dd7..5f42fe35eef 100644 --- a/test/integration/scheduler/taint_test.go +++ b/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