From 1158040a8ed38a6c11dd9caccb55a1df8be479ac Mon Sep 17 00:00:00 2001 From: shaloulcy Date: Fri, 26 Jan 2024 15:11:08 +0800 Subject: [PATCH] webhook: add node affinity for pod when MultiQuotaTree on (#1864) Signed-off-by: chuanyun.lcy Co-authored-by: chuanyun.lcy --- pkg/util/fieldindex/register.go | 12 + .../pod/mutating/multi_quota_tree_affinity.go | 149 ++++++ .../multi_quota_tree_affinity_test.go | 436 ++++++++++++++++++ pkg/webhook/pod/mutating/mutating_handler.go | 5 + 4 files changed, 602 insertions(+) create mode 100644 pkg/webhook/pod/mutating/multi_quota_tree_affinity.go create mode 100644 pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go diff --git a/pkg/util/fieldindex/register.go b/pkg/util/fieldindex/register.go index 9b1e1bef6..c29964210 100644 --- a/pkg/util/fieldindex/register.go +++ b/pkg/util/fieldindex/register.go @@ -85,6 +85,18 @@ var indexDescriptors = []fieldIndexDescriptor{ return extension.GetAnnotationQuotaNamespaces(eq) }, }, + { + description: "index elastic quota by name", + obj: &apiv1alpha1.ElasticQuota{}, + field: "metadata.name", + indexerFunc: func(obj client.Object) []string { + eq, ok := obj.(*apiv1alpha1.ElasticQuota) + if !ok { + return []string{} + } + return []string{eq.Name} + }, + }, } func RegisterFieldIndexes(c cache.Cache) error { diff --git a/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go b/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go new file mode 100644 index 000000000..1591bca90 --- /dev/null +++ b/pkg/webhook/pod/mutating/multi_quota_tree_affinity.go @@ -0,0 +1,149 @@ +/* +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 mutating + +import ( + "context" + + 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" +) + +func (h *PodMutatingHandler) addNodeAffinityForMultiQuotaTree(ctx context.Context, req admission.Request, pod *corev1.Pod) error { + if req.Operation != admissionv1.Create { + return nil + } + + if !utilfeature.DefaultFeatureGate.Enabled(features.MultiQuotaTree) { + return nil + } + + 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] + } + + treeID := extension.GetQuotaTreeID(quota) + if treeID == "" { + return nil + } + + profileList := "av1alpha1.ElasticQuotaProfileList{} + err := h.Client.List(ctx, profileList, &client.ListOptions{ + LabelSelector: labels.SelectorFromSet(map[string]string{extension.LabelQuotaTreeID: treeID}), + }, utilclient.DisableDeepCopy) + if err != nil { + return err + } + + if len(profileList.Items) == 0 { + return nil + } + + nodeSelector := profileList.Items[0].Spec.NodeSelector + if nodeSelector == nil { + return nil + } + + requirements := convertNodeSelectorToNodeSelectorRequirements(nodeSelector) + + affinity := pod.Spec.Affinity + if affinity == nil { + affinity = &corev1.Affinity{} + pod.Spec.Affinity = affinity + } + + nodeAffinity := affinity.NodeAffinity + if nodeAffinity == nil { + nodeAffinity = &corev1.NodeAffinity{} + affinity.NodeAffinity = nodeAffinity + } + + required := nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution + if required == nil { + required = &corev1.NodeSelector{} + nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = required + } + + for i, term := range required.NodeSelectorTerms { + term.MatchExpressions = append(term.MatchExpressions, requirements...) + required.NodeSelectorTerms[i] = term + } + + if len(required.NodeSelectorTerms) == 0 { + required.NodeSelectorTerms = []corev1.NodeSelectorTerm{ + { + MatchExpressions: requirements, + }, + } + } + + return nil +} + +func convertNodeSelectorToNodeSelectorRequirements(nodeSelector *metav1.LabelSelector) []corev1.NodeSelectorRequirement { + requirements := make([]corev1.NodeSelectorRequirement, 0, len(nodeSelector.MatchLabels)+len(nodeSelector.MatchExpressions)) + for k, v := range nodeSelector.MatchLabels { + requirements = append(requirements, corev1.NodeSelectorRequirement{ + Key: k, + Operator: corev1.NodeSelectorOpIn, + Values: []string{v}, + }) + } + for _, expression := range nodeSelector.MatchExpressions { + requirements = append(requirements, corev1.NodeSelectorRequirement{ + Key: expression.Key, + Operator: corev1.NodeSelectorOperator(string(expression.Operator)), + Values: expression.Values, + }) + } + + return requirements +} diff --git a/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go b/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go new file mode 100644 index 000000000..5d555a805 --- /dev/null +++ b/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go @@ -0,0 +1,436 @@ +/* +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 mutating + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + admissionv1 "k8s.io/api/admission/v1" + 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" + "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" + quotav1alpha1 "github.com/koordinator-sh/koordinator/apis/quota/v1alpha1" + "github.com/koordinator-sh/koordinator/pkg/features" + "github.com/koordinator-sh/koordinator/pkg/util/feature" +) + +func init() { + _ = quotav1alpha1.AddToScheme(scheme.Scheme) + _ = schedulingv1alpha1.AddToScheme(scheme.Scheme) +} + +func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) { + testCases := []struct { + name string + pod *corev1.Pod + quotas []*schedulingv1alpha1.ElasticQuota + profile *quotav1alpha1.ElasticQuotaProfile + expected *corev1.Pod + }{ + { + name: "no quota label", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + }, + Spec: corev1.PodSpec{}, + }, + quotas: nil, + profile: nil, + expected: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + }, + Spec: corev1.PodSpec{}, + }, + }, + { + name: "no quota profile", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + Labels: map[string]string{ + extension.LabelQuotaName: "quota1", + }, + }, + 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", + }, + }, + Spec: corev1.PodSpec{}, + }, + }, + { + name: "add node affinity", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + Labels: map[string]string{ + extension.LabelQuotaName: "quota1", + }, + }, + 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", + }, + }, + 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{"test"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "append node affinity", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + Labels: map[string]string{ + extension.LabelQuotaName: "quota1", + }, + }, + Spec: corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + { + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: "topology.kubernetes.io/zone", + Operator: corev1.NodeSelectorOpIn, + Values: []string{"cn-hangzhou-a"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + 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", + }, + }, + Spec: corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + { + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: "topology.kubernetes.io/zone", + Operator: corev1.NodeSelectorOpIn, + Values: []string{"cn-hangzhou-a"}, + }, + { + Key: "node-pool", + Operator: corev1.NodeSelectorOpIn, + Values: []string{"test"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "multi quotas", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + Labels: map[string]string{ + extension.LabelQuotaName: "quota2", + }, + }, + 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", + }, + }, + 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{"test"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "default quota", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "test-pod-1", + }, + Spec: corev1.PodSpec{}, + }, + quotas: []*schedulingv1alpha1.ElasticQuota{ + { + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "default", + 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", + }, + 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{"test"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + } + + 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.Equal(tc.expected, tc.pod) + }) + } +} diff --git a/pkg/webhook/pod/mutating/mutating_handler.go b/pkg/webhook/pod/mutating/mutating_handler.go index 88392c1d4..1c8874bd2 100644 --- a/pkg/webhook/pod/mutating/mutating_handler.go +++ b/pkg/webhook/pod/mutating/mutating_handler.go @@ -108,6 +108,11 @@ func (h *PodMutatingHandler) handleCreate(ctx context.Context, req admission.Req return err } + if err := h.addNodeAffinityForMultiQuotaTree(ctx, req, obj); err != nil { + klog.Errorf("Failed to mutating Pod %s/%s by MultiQuotaTree, err: %v", obj.Namespace, obj.Name, err) + return err + } + return nil }