1/*
2Copyright 2019 The Kubernetes Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17package testing
18
19import (
20	"k8s.io/apimachinery/pkg/runtime/schema"
21	"k8s.io/kube-scheduler/config/v1beta2"
22	schedulerapi "k8s.io/kubernetes/pkg/scheduler/apis/config"
23	"k8s.io/kubernetes/pkg/scheduler/apis/config/scheme"
24	"k8s.io/kubernetes/pkg/scheduler/framework"
25	"k8s.io/kubernetes/pkg/scheduler/framework/runtime"
26)
27
28var configDecoder = scheme.Codecs.UniversalDecoder()
29
30// NewFramework creates a Framework from the register functions and options.
31func NewFramework(fns []RegisterPluginFunc, profileName string, opts ...runtime.Option) (framework.Framework, error) {
32	registry := runtime.Registry{}
33	profile := &schedulerapi.KubeSchedulerProfile{
34		SchedulerName: profileName,
35		Plugins:       &schedulerapi.Plugins{},
36	}
37	for _, f := range fns {
38		f(&registry, profile)
39	}
40	return runtime.NewFramework(registry, profile, opts...)
41}
42
43// RegisterPluginFunc is a function signature used in method RegisterFilterPlugin()
44// to register a Filter Plugin to a given registry.
45type RegisterPluginFunc func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile)
46
47// RegisterQueueSortPlugin returns a function to register a QueueSort Plugin to a given registry.
48func RegisterQueueSortPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
49	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "QueueSort")
50}
51
52// RegisterPreFilterPlugin returns a function to register a PreFilter Plugin to a given registry.
53func RegisterPreFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
54	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreFilter")
55}
56
57// RegisterFilterPlugin returns a function to register a Filter Plugin to a given registry.
58func RegisterFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
59	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Filter")
60}
61
62// RegisterReservePlugin returns a function to register a Reserve Plugin to a given registry.
63func RegisterReservePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
64	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Reserve")
65}
66
67// RegisterPermitPlugin returns a function to register a Permit Plugin to a given registry.
68func RegisterPermitPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
69	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Permit")
70}
71
72// RegisterPreBindPlugin returns a function to register a PreBind Plugin to a given registry.
73func RegisterPreBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
74	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreBind")
75}
76
77// RegisterScorePlugin returns a function to register a Score Plugin to a given registry.
78func RegisterScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory, weight int32) RegisterPluginFunc {
79	return RegisterPluginAsExtensionsWithWeight(pluginName, weight, pluginNewFunc, "Score")
80}
81
82// RegisterPreScorePlugin returns a function to register a Score Plugin to a given registry.
83func RegisterPreScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
84	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreScore")
85}
86
87// RegisterBindPlugin returns a function to register a Bind Plugin to a given registry.
88func RegisterBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
89	return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Bind")
90}
91
92// RegisterPluginAsExtensions returns a function to register a Plugin as given extensionPoints to a given registry.
93func RegisterPluginAsExtensions(pluginName string, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc {
94	return RegisterPluginAsExtensionsWithWeight(pluginName, 1, pluginNewFunc, extensions...)
95}
96
97// RegisterPluginAsExtensionsWithWeight returns a function to register a Plugin as given extensionPoints with weight to a given registry.
98func RegisterPluginAsExtensionsWithWeight(pluginName string, weight int32, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc {
99	return func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile) {
100		reg.Register(pluginName, pluginNewFunc)
101		for _, extension := range extensions {
102			ps := getPluginSetByExtension(profile.Plugins, extension)
103			if ps == nil {
104				continue
105			}
106			ps.Enabled = append(ps.Enabled, schedulerapi.Plugin{Name: pluginName, Weight: weight})
107		}
108		// Use defaults from latest config API version.
109		var gvk schema.GroupVersionKind
110		gvk = v1beta2.SchemeGroupVersion.WithKind(pluginName + "Args")
111		if args, _, err := configDecoder.Decode(nil, &gvk, nil); err == nil {
112			profile.PluginConfig = append(profile.PluginConfig, schedulerapi.PluginConfig{
113				Name: pluginName,
114				Args: args,
115			})
116		}
117	}
118}
119
120func getPluginSetByExtension(plugins *schedulerapi.Plugins, extension string) *schedulerapi.PluginSet {
121	switch extension {
122	case "QueueSort":
123		return &plugins.QueueSort
124	case "Filter":
125		return &plugins.Filter
126	case "PreFilter":
127		return &plugins.PreFilter
128	case "PreScore":
129		return &plugins.PreScore
130	case "Score":
131		return &plugins.Score
132	case "Bind":
133		return &plugins.Bind
134	case "Reserve":
135		return &plugins.Reserve
136	case "Permit":
137		return &plugins.Permit
138	case "PreBind":
139		return &plugins.PreBind
140	case "PostBind":
141		return &plugins.PostBind
142	default:
143		return nil
144	}
145}
146