webhook: improve node affinity performance for MultiQuotaTree

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)
+}

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)
+}

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)
 }

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)
+}

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
+}

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 = &quotaList.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 := &quotav1alpha1.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