diff --git a/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go b/pkg/webhook/elasticquota/plugin_check_quota_meta_validate.go index 9bb0fd6c7..0b657d134 100644 --- 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 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..1690eafaa --- /dev/null +++ b/pkg/webhook/elasticquota/quota_handler_test.go @@ -0,0 +1,59 @@ +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 4099daa60..d381c4079 100644 --- 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 index a79e64548..4c4719cd8 100644 --- 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 index 24b6f6089..081183c7e 100644 --- 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 {