webhook: fix OnQuotaAdd/Update/Delete when quota namespace annotation…

… changed Signed-off-by: chuanyun.lcy <chuanyun.lcy@alibaba-inc.com>
koordinator-sh · Jan 29, 2024 · 77f2b17 · 77f2b17
1 parent cb3703b
commit 77f2b17
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 5 deletions.
diff --git a/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go
@@ -115,10 +115,10 @@ func (c *QuotaMetaChecker) GetQuotaTopologyInfo() *QuotaTopologySummary {
 	return quotaMetaCheck.QuotaTopo.getQuotaTopologyInfo()
 }
 
-func (c *QuotaMetaChecker) GetQuotaInfo(name string) *QuotaInfo {
+func (c *QuotaMetaChecker) GetQuotaInfo(name, namespace string) *QuotaInfo {
 	if c.QuotaTopo == nil {
 		return nil
 	}
 
-	return c.QuotaTopo.getQuotaInfo(name)
+	return c.QuotaTopo.getQuotaInfo(name, namespace)
 }
diff --git 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_topology.go b/pkg/webhook/elasticquota/quota_topology.go
@@ -264,6 +264,17 @@ func (qt *quotaTopology) getQuotaTopologyInfo() *QuotaTopologySummary {
 	return result
 }
 
-func (qt *quotaTopology) getQuotaInfo(name string) *QuotaInfo {
-	return qt.quotaInfoMap[name]
+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
@@ -49,7 +49,7 @@ func (h *PodMutatingHandler) addNodeAffinityForMultiQuotaTree(ctx context.Contex
 		quotaName = pod.Namespace
 	}
 
-	info := plugin.GetQuotaInfo(quotaName)
+	info := plugin.GetQuotaInfo(quotaName, pod.Namespace)
 	if info == nil {
 		return nil
 	}

diff --git a/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go b/pkg/webhook/pod/mutating/multi_quota_tree_affinity_test.go
@@ -132,6 +132,9 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) {
 					extension.LabelQuotaTreeID:   "tree2",
 					extension.LabelQuotaIsParent: "true",
 				},
+				Annotations: map[string]string{
+					extension.AnnotationQuotaNamespaces: "[\"namespace2\"]",
+				},
 			},
 		},
 		// the children quotas of root-quota-b
@@ -388,6 +391,41 @@ func TestAddNodeAffinityForMultiQuotaTree(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "default quota 2",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "namespace2",
+					Name:      "test-pod-1",
+				},
+				Spec: corev1.PodSpec{},
+			},
+			expected: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "namespace2",
+					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{"nodePoolB"},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {