diff --git a/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go index 85b6454f6..0b657d134 100644 --- a/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go +++ b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go @@ -114,3 +114,11 @@ func (c *QuotaMetaChecker) GetQuotaTopologyInfo() *QuotaTopologySummary { } return quotaMetaCheck.QuotaTopo.getQuotaTopologyInfo() } + +func (c *QuotaMetaChecker) GetQuotaInfo(name, namespace string) *QuotaInfo { + if c.QuotaTopo == nil { + return nil + } + + return c.QuotaTopo.getQuotaInfo(name, namespace) +} diff --git a/pkg/webhook/elasticquota/plugin_check_quota_meta_validate_test.go b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate_test.go new file mode 100644 index 000000000..b3ff7ac11 --- /dev/null +++ b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate_test.go @@ -0,0 +1,70 @@ +/* +Copyright 2022 The Koordinator Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package elasticquota + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + admissionv1 "k8s.io/api/admission/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/webhook/admission" + "sigs.k8s.io/scheduler-plugins/pkg/apis/scheduling/v1alpha1" + + "github.com/koordinator-sh/koordinator/apis/extension" +) + +func TestQuotaMetaChecker(t *testing.T) { + client := fake.NewClientBuilder().Build() + sche := client.Scheme() + sche.AddKnownTypes(schema.GroupVersion{ + Group: "scheduling.sigs.k8s.io", + Version: "v1alpha1", + }, &v1alpha1.ElasticQuota{}, &v1alpha1.ElasticQuotaList{}) + decoder, _ := admission.NewDecoder(sche) + + plugin := NewPlugin(decoder, client) + + request := admission.Request{ + AdmissionRequest: admissionv1.AdmissionRequest{ + Resource: metav1.GroupVersionResource{ + Group: "scheduling.sigs.k8s.io", + Version: "v1alpha1", + Resource: "elasticquotas", + }, + Operation: admissionv1.Create, + Object: runtime.RawExtension{}, + }, + } + + parentQuota := MakeQuota("parentQuota").Namespace("kube-system").Max(MakeResourceList().CPU(120).Mem(1048576).Obj()). + Min(MakeResourceList().CPU(120).Mem(1048576).Obj()).IsParent(true).Obj() + + // validate quota + err := plugin.ValidateQuota(context.TODO(), request, parentQuota) + assert.Nil(t, err) + + // get quota info + quotaInfo := plugin.GetQuotaInfo(parentQuota.Name, parentQuota.Namespace) + assert.NotNil(t, quotaInfo) + assert.Equal(t, parentQuota.Name, quotaInfo.Name) + assert.Equal(t, extension.RootQuotaName, quotaInfo.ParentName) +} diff --git a/pkg/webhook/elasticquota/quota_handler.go b/pkg/webhook/elasticquota/quota_handler.go index a2f6e6ca9..767cbb8fb 100644 --- a/pkg/webhook/elasticquota/quota_handler.go +++ b/pkg/webhook/elasticquota/quota_handler.go @@ -17,6 +17,8 @@ limitations under the License. package elasticquota import ( + "reflect" + "k8s.io/klog/v2" "github.com/koordinator-sh/koordinator/apis/extension" @@ -44,6 +46,12 @@ func (qt *quotaTopology) OnQuotaAdd(obj interface{}) { qt.quotaHierarchyInfo[quotaInfo.ParentName] = make(map[string]struct{}) } qt.quotaHierarchyInfo[quotaInfo.ParentName][quotaInfo.Name] = struct{}{} + + namespaces := extension.GetAnnotationQuotaNamespaces(quota) + for _, ns := range namespaces { + qt.namespaceToQuotaMap[ns] = quota.Name + } + klog.V(5).Infof("OnQuotaAdd success: %v.%v", quota.Namespace, quota.Name) } @@ -71,6 +79,18 @@ func (qt *quotaTopology) OnQuotaUpdate(oldObj, newObj interface{}) { delete(qt.quotaHierarchyInfo[oldQuotaInfo.ParentName], oldQuotaInfo.Name) qt.quotaHierarchyInfo[newQuotaInfo.ParentName][newQuotaInfo.Name] = struct{}{} } + + oldNamespaces := extension.GetAnnotationQuotaNamespaces(oldQuota) + newNamespaces := extension.GetAnnotationQuotaNamespaces(newQuota) + if !reflect.DeepEqual(oldNamespaces, newNamespaces) { + for _, ns := range oldNamespaces { + delete(qt.namespaceToQuotaMap, ns) + } + for _, ns := range newNamespaces { + qt.namespaceToQuotaMap[ns] = newQuota.Name + } + } + klog.V(5).Infof("OnQuotaUpdate success: %v.%v", newQuota.Namespace, newQuota.Name) } @@ -89,5 +109,10 @@ func (qt *quotaTopology) OnQuotaDelete(obj interface{}) { delete(qt.quotaHierarchyInfo[parentName], quota.Name) delete(qt.quotaHierarchyInfo, quota.Name) delete(qt.quotaInfoMap, quota.Name) + + namespaces := extension.GetAnnotationQuotaNamespaces(quota) + for _, ns := range namespaces { + delete(qt.namespaceToQuotaMap, ns) + } klog.V(5).Infof("OnQuotaDelete success: %v.%v", quota.Namespace, quota.Name) } diff --git a/pkg/webhook/elasticquota/quota_handler_test.go b/pkg/webhook/elasticquota/quota_handler_test.go new file mode 100644 index 000000000..a42b91cbc --- /dev/null +++ b/pkg/webhook/elasticquota/quota_handler_test.go @@ -0,0 +1,75 @@ +/* +Copyright 2022 The Koordinator Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package elasticquota + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + "github.com/koordinator-sh/koordinator/apis/extension" +) + +func TestQuotaHandler(t *testing.T) { + client := fake.NewClientBuilder().Build() + topology := NewQuotaTopology(client) + + parentQuota := MakeQuota("parentQuota").Namespace("kube-system").Max(MakeResourceList().CPU(120).Mem(1048576).Obj()). + Min(MakeResourceList().CPU(120).Mem(1048576).Obj()).IsParent(true).Obj() + + topology.OnQuotaAdd(parentQuota) + + // check parent quota info + quotaInfo := topology.getQuotaInfo(parentQuota.Name, parentQuota.Namespace) + assert.NotNil(t, quotaInfo) + assert.Equal(t, parentQuota.Name, quotaInfo.Name) + assert.Equal(t, extension.RootQuotaName, quotaInfo.ParentName) + + childQuota := MakeQuota("childQuota").Namespace("kube-system").Max(MakeResourceList().CPU(120).Mem(1048576).Obj()). + Min(MakeResourceList().CPU(120).Mem(1048576).Obj()).IsParent(false).ParentName(parentQuota.Name).Annotations( + map[string]string{extension.AnnotationQuotaNamespaces: `["namespace1"]`}, + ).Obj() + topology.OnQuotaAdd(childQuota) + + // check child quota info + quotaInfo = topology.getQuotaInfo(childQuota.Name, childQuota.Namespace) + assert.NotNil(t, quotaInfo) + assert.Equal(t, childQuota.Name, quotaInfo.Name) + assert.Equal(t, parentQuota.Name, quotaInfo.ParentName) + // get quota by namespace + quotaInfo = topology.getQuotaInfo("", "namespace1") + assert.NotNil(t, quotaInfo) + assert.Equal(t, childQuota.Name, quotaInfo.Name) + + // update quota info + newChildQuota := childQuota.DeepCopy() + newChildQuota.Annotations[extension.AnnotationQuotaNamespaces] = `["namespace2"]` + topology.OnQuotaUpdate(childQuota, newChildQuota) + quotaInfo = topology.getQuotaInfo("", "namespace1") + assert.Nil(t, quotaInfo) + quotaInfo = topology.getQuotaInfo("", "namespace2") + assert.NotNil(t, quotaInfo) + assert.Equal(t, childQuota.Name, quotaInfo.Name) + + // delete quota + topology.OnQuotaDelete(newChildQuota) + quotaInfo = topology.getQuotaInfo(childQuota.Name, childQuota.Namespace) + assert.Nil(t, quotaInfo) + quotaInfo = topology.getQuotaInfo("", "namespace2") + assert.Nil(t, quotaInfo) +} diff --git a/pkg/webhook/elasticquota/quota_topology.go b/pkg/webhook/elasticquota/quota_topology.go index 6e301b666..d381c4079 100644 --- a/pkg/webhook/elasticquota/quota_topology.go +++ b/pkg/webhook/elasticquota/quota_topology.go @@ -263,3 +263,18 @@ func (qt *quotaTopology) getQuotaTopologyInfo() *QuotaTopologySummary { } return result } + +func (qt *quotaTopology) getQuotaInfo(name, namespace string) *QuotaInfo { + qt.lock.Lock() + defer qt.lock.Unlock() + + info, ok := qt.quotaInfoMap[name] + if ok { + return info + } + quotaName, ok := qt.namespaceToQuotaMap[namespace] + if ok { + return qt.quotaInfoMap[quotaName] + } + return nil +} diff --git a/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go b/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go index 1591bca90..4c4719cd8 100644 --- a/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go +++ b/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go @@ -21,20 +21,17 @@ import ( admissionv1 "k8s.io/api/admission/v1" corev1 "k8s.io/api/core/v1" - "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/fields" "k8s.io/apimachinery/pkg/labels" - "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/webhook/admission" - schedulingv1alpha1 "sigs.k8s.io/scheduler-plugins/pkg/apis/scheduling/v1alpha1" "github.com/koordinator-sh/koordinator/apis/extension" quotav1alpha1 "github.com/koordinator-sh/koordinator/apis/quota/v1alpha1" "github.com/koordinator-sh/koordinator/pkg/features" utilclient "github.com/koordinator-sh/koordinator/pkg/util/client" utilfeature "github.com/koordinator-sh/koordinator/pkg/util/feature" + "github.com/koordinator-sh/koordinator/pkg/webhook/elasticquota" ) func (h *PodMutatingHandler) addNodeAffinityForMultiQuotaTree(ctx context.Context, req admission.Request, pod *corev1.Pod) error { @@ -46,38 +43,23 @@ func (h *PodMutatingHandler) addNodeAffinityForMultiQuotaTree(ctx context.Contex return nil } + plugin := elasticquota.NewPlugin(h.Decoder, h.Client) quotaName := extension.GetQuotaName(pod) - quota := &schedulingv1alpha1.ElasticQuota{} if quotaName == "" { - err := h.Client.Get(ctx, types.NamespacedName{Namespace: pod.Namespace, Name: pod.Namespace}, quota) - if err != nil { - if errors.IsNotFound(err) { - return nil - } - return err - } - } else { - quotaList := &schedulingv1alpha1.ElasticQuotaList{} - err := h.Client.List(ctx, quotaList, &client.ListOptions{ - FieldSelector: fields.OneTermEqualSelector("metadata.name", quotaName), - }, utilclient.DisableDeepCopy) - if err != nil { - return err - } - if len(quotaList.Items) == 0 { - return nil - } - quota = "aList.Items[0] + quotaName = pod.Namespace } - treeID := extension.GetQuotaTreeID(quota) - if treeID == "" { + info := plugin.GetQuotaInfo(quotaName, pod.Namespace) + if info == nil { + return nil + } + if info.TreeID == "" { return nil } profileList := "av1alpha1.ElasticQuotaProfileList{} err := h.Client.List(ctx, profileList, &client.ListOptions{ - LabelSelector: labels.SelectorFromSet(map[string]string{extension.LabelQuotaTreeID: treeID}), + LabelSelector: labels.SelectorFromSet(map[string]string{extension.LabelQuotaTreeID: info.TreeID}), }, utilclient.DisableDeepCopy) if err != nil { return err diff --git a/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go b/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go index 5d555a805..081183c7e 100644 --- a/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go +++ b/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go @@ -25,9 +25,8 @@ import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/kubernetes/scheme" - "sigs.k8s.io/controller-runtime/pkg/client/fake" - "sigs.k8s.io/controller-runtime/pkg/webhook/admission" schedulingv1alpha1 "sigs.k8s.io/scheduler-plugins/pkg/apis/scheduling/v1alpha1" "github.com/koordinator-sh/koordinator/apis/extension" @@ -42,11 +41,133 @@ func init() { } func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { + handler, fakeInformers := makeTestHandler() + quotaInformer, err := fakeInformers.FakeInformerFor(&schedulingv1alpha1.ElasticQuota{}) + assert.NoError(t, err) + + profiles := []*quotav1alpha1.ElasticQuotaProfile{ + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "profileA", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree1", + }, + }, + Spec: quotav1alpha1.ElasticQuotaProfileSpec{ + QuotaName: "root-quota-a", + NodeSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "node-pool": "nodePoolA", + }, + }, + }, + }, + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "profileB", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree2", + }, + }, + Spec: quotav1alpha1.ElasticQuotaProfileSpec{ + QuotaName: "root-quota-b", + NodeSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "node-pool": "nodePoolB", + }, + }, + }, + }, + } + + quotas := []*schedulingv1alpha1.ElasticQuota{ + // other-quota, no tree + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "other-quota", + }, + }, + // root-quota-a + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "root-quota-a", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree1", + extension.LabelQuotaIsParent: "true", + }, + }, + }, + // the children quotas of root-quota-a + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "root-quota-a-child1", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree1", + extension.LabelQuotaParent: "root-quota-a", + }, + }, + }, + // the namespace quota + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "namespace1", + Name: "namespace1", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree1", + extension.LabelQuotaParent: "root-quota-a", + }, + }, + }, + // root-quota-b + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "root-quota-b", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree2", + extension.LabelQuotaIsParent: "true", + }, + Annotations: map[string]string{ + extension.AnnotationQuotaNamespaces: "[\"namespace2\"]", + }, + }, + }, + // the children quotas of root-quota-b + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "kube-system", + Name: "root-quota-b-child1", + Labels: map[string]string{ + extension.LabelQuotaTreeID: "tree2", + extension.LabelQuotaParent: "root-quota-b", + }, + }, + }, + } + + for _, profile := range profiles { + err := handler.Client.Create(context.TODO(), profile) + assert.NoError(t, err) + } + + for _, quota := range quotas { + err := handler.Client.Create(context.TODO(), quota) + assert.NoError(t, err) + quotaInformer.Add(quota) + } + + quota := &schedulingv1alpha1.ElasticQuota{} + err = handler.Client.Get(context.TODO(), types.NamespacedName{Namespace: "kube-system", Name: "root-quota-b-child1"}, quota) + assert.NoError(t, err) + testCases := []struct { name string pod *corev1.Pod - quotas []*schedulingv1alpha1.ElasticQuota - profile *quotav1alpha1.ElasticQuotaProfile expected *corev1.Pod }{ { @@ -58,8 +179,6 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { }, Spec: corev1.PodSpec{}, }, - quotas: nil, - profile: nil, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "default", @@ -75,26 +194,17 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "other-quota", }, }, Spec: corev1.PodSpec{}, }, - quotas: []*schedulingv1alpha1.ElasticQuota{ - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1", - }, - }, - }, - profile: nil, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "other-quota", }, }, Spec: corev1.PodSpec{}, @@ -107,45 +217,17 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "root-quota-a-child1", }, }, Spec: corev1.PodSpec{}, }, - quotas: []*schedulingv1alpha1.ElasticQuota{ - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - }, - }, - profile: "av1alpha1.ElasticQuotaProfile{ - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1-profile", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - Spec: quotav1alpha1.ElasticQuotaProfileSpec{ - QuotaName: "root-quota", - NodeSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "node-pool": "test", - }, - }, - }, - }, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "root-quota-a-child1", }, }, Spec: corev1.PodSpec{ @@ -158,7 +240,7 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { { Key: "node-pool", Operator: corev1.NodeSelectorOpIn, - Values: []string{"test"}, + Values: []string{"nodePoolA"}, }, }, }, @@ -176,7 +258,7 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "root-quota-a-child1", }, }, Spec: corev1.PodSpec{ @@ -199,40 +281,12 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { }, }, }, - quotas: []*schedulingv1alpha1.ElasticQuota{ - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - }, - }, - profile: "av1alpha1.ElasticQuotaProfile{ - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1-profile", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - Spec: quotav1alpha1.ElasticQuotaProfileSpec{ - QuotaName: "root-quota", - NodeSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "node-pool": "test", - }, - }, - }, - }, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota1", + extension.LabelQuotaName: "root-quota-a-child1", }, }, Spec: corev1.PodSpec{ @@ -250,7 +304,7 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { { Key: "node-pool", Operator: corev1.NodeSelectorOpIn, - Values: []string{"test"}, + Values: []string{"nodePoolA"}, }, }, }, @@ -268,54 +322,17 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota2", + extension.LabelQuotaName: "root-quota-b-child1", }, }, Spec: corev1.PodSpec{}, }, - quotas: []*schedulingv1alpha1.ElasticQuota{ - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - }, - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota2", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - }, - }, - profile: "av1alpha1.ElasticQuotaProfile{ - ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1-profile", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - Spec: quotav1alpha1.ElasticQuotaProfileSpec{ - QuotaName: "root-quota", - NodeSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "node-pool": "test", - }, - }, - }, - }, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "default", Name: "test-pod-1", Labels: map[string]string{ - extension.LabelQuotaName: "quota2", + extension.LabelQuotaName: "root-quota-b-child1", }, }, Spec: corev1.PodSpec{ @@ -328,7 +345,7 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { { Key: "node-pool", Operator: corev1.NodeSelectorOpIn, - Values: []string{"test"}, + Values: []string{"nodePoolB"}, }, }, }, @@ -343,42 +360,49 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { name: "default quota", pod: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ - Namespace: "default", + Namespace: "namespace1", Name: "test-pod-1", }, Spec: corev1.PodSpec{}, }, - quotas: []*schedulingv1alpha1.ElasticQuota{ - { - ObjectMeta: metav1.ObjectMeta{ - Namespace: "default", - Name: "default", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", + expected: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "namespace1", + Name: "test-pod-1", + }, + Spec: corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + { + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: "node-pool", + Operator: corev1.NodeSelectorOpIn, + Values: []string{"nodePoolA"}, + }, + }, + }, + }, + }, }, }, }, }, - profile: "av1alpha1.ElasticQuotaProfile{ + }, + { + name: "default quota 2", + pod: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ - Namespace: "kube-system", - Name: "quota1-profile", - Labels: map[string]string{ - extension.LabelQuotaTreeID: "123456", - }, - }, - Spec: quotav1alpha1.ElasticQuotaProfileSpec{ - QuotaName: "root-quota", - NodeSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "node-pool": "test", - }, - }, + Namespace: "namespace2", + Name: "test-pod-1", }, + Spec: corev1.PodSpec{}, }, expected: &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ - Namespace: "default", + Namespace: "namespace2", Name: "test-pod-1", }, Spec: corev1.PodSpec{ @@ -391,7 +415,7 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { { Key: "node-pool", Operator: corev1.NodeSelectorOpIn, - Values: []string{"test"}, + Values: []string{"nodePoolB"}, }, }, }, @@ -407,30 +431,12 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { defer feature.SetFeatureGateDuringTest(t, feature.DefaultMutableFeatureGate, features.MultiQuotaTree, true)() - assert := assert.New(t) - - client := fake.NewClientBuilder().Build() - decoder, _ := admission.NewDecoder(scheme.Scheme) - handler := &PodMutatingHandler{ - Client: client, - Decoder: decoder, - } - - if tc.profile != nil { - err := client.Create(context.TODO(), tc.profile) - assert.NoError(err) - } - - for _, quota := range tc.quotas { - err := client.Create(context.TODO(), quota) - assert.NoError(err) - } req := newAdmission(admissionv1.Create, runtime.RawExtension{}, runtime.RawExtension{}, "") err := handler.addNodeAffinityForMultiQuotaTree(context.TODO(), req, tc.pod) - assert.NoError(err) + assert.NoError(t, err) - assert.Equal(tc.expected, tc.pod) + assert.Equal(t, tc.expected, tc.pod) }) } } diff --git a/pkg/webhook/pod/mutating/mutating_handler_test.go b/pkg/webhook/pod/mutating/mutating_handler_test.go index c5a372404..0e93bc9a1 100644 --- a/pkg/webhook/pod/mutating/mutating_handler_test.go +++ b/pkg/webhook/pod/mutating/mutating_handler_test.go @@ -25,18 +25,37 @@ import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "k8s.io/client-go/kubernetes/scheme" + "k8s.io/apimachinery/pkg/runtime/schema" + clientcache "k8s.io/client-go/tools/cache" + sigcache "sigs.k8s.io/controller-runtime/pkg/cache" + "sigs.k8s.io/controller-runtime/pkg/cache/informertest" "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/runtime/inject" "sigs.k8s.io/controller-runtime/pkg/webhook/admission" + "sigs.k8s.io/scheduler-plugins/pkg/apis/scheduling/v1alpha1" + + "github.com/koordinator-sh/koordinator/pkg/webhook/elasticquota" ) -func makeTestHandler() *PodMutatingHandler { +func makeTestHandler() (*PodMutatingHandler, *informertest.FakeInformers) { client := fake.NewClientBuilder().Build() - decoder, _ := admission.NewDecoder(scheme.Scheme) + sche := client.Scheme() + sche.AddKnownTypes(schema.GroupVersion{ + Group: "scheduling.sigs.k8s.io", + Version: "v1alpha1", + }, &v1alpha1.ElasticQuota{}, &v1alpha1.ElasticQuotaList{}) + decoder, _ := admission.NewDecoder(sche) handler := &PodMutatingHandler{} handler.InjectClient(client) handler.InjectDecoder(decoder) - return handler + + cacheTmp := &informertest.FakeInformers{ + InformersByGVK: map[schema.GroupVersionKind]clientcache.SharedIndexInformer{}, + Scheme: sche, + } + handler.InjectCache(cacheTmp) + + return handler, cacheTmp } func gvr(resource string) metav1.GroupVersionResource { @@ -48,7 +67,7 @@ func gvr(resource string) metav1.GroupVersionResource { } func TestMutatingHandler(t *testing.T) { - handler := makeTestHandler() + handler, _ := makeTestHandler() ctx := context.Background() testCases := []struct { @@ -117,3 +136,21 @@ func TestMutatingHandler(t *testing.T) { }) } } + +var _ inject.Cache = &PodMutatingHandler{} + +func (h *PodMutatingHandler) InjectCache(cache sigcache.Cache) error { + ctx := context.TODO() + quotaInformer, err := cache.GetInformer(ctx, &v1alpha1.ElasticQuota{}) + if err != nil { + return err + } + plugin := elasticquota.NewPlugin(h.Decoder, h.Client) + qt := plugin.QuotaTopo + quotaInformer.AddEventHandler(clientcache.ResourceEventHandlerFuncs{ + AddFunc: qt.OnQuotaAdd, + UpdateFunc: qt.OnQuotaUpdate, + DeleteFunc: qt.OnQuotaDelete, + }) + return nil +}