fix GetPodEffectiveDecorations and add ut

pkg/controllers/collaset/collaset_controller.go

@@ -39,7 +39,6 @@ import (
 	collasetutils "kusionstack.io/operating/pkg/controllers/collaset/utils"
 	controllerutils "kusionstack.io/operating/pkg/controllers/utils"
 	"kusionstack.io/operating/pkg/controllers/utils/expectations"
-	utilspoddecoration "kusionstack.io/operating/pkg/controllers/utils/poddecoration"
 	"kusionstack.io/operating/pkg/controllers/utils/podopslifecycle"
 	"kusionstack.io/operating/pkg/controllers/utils/revision"
 	commonutils "kusionstack.io/operating/pkg/utils"
@@ -173,7 +172,7 @@ func (r *CollaSetReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		UpdatedRevision: updatedRevision,
 		NewStatus:       newStatus,
 	}
-	resources.PodDecorations, resources.OldRevisionDecorations, err = utilspoddecoration.GetEffectiveDecorationsByCollaSet(ctx, r.Client, instance)
+	resources.PodDecorations, resources.OldRevisionDecorations, err = utils.GetEffectiveDecorationsByCollaSet(ctx, r.Client, instance)
 	if err != nil {
 		return ctrl.Result{}, fmt.Errorf("fail to get effective pod decorations by CollaSet %s: %s", key, err)
 	}

pkg/controllers/collaset/utils/poddecoration.go

@@ -0,0 +1,72 @@
+/*
+Copyright 2023 The KusionStack 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 utils
+
+import (
+	"context"
+	"fmt"
+
+	appsv1 "k8s.io/api/apps/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	appsv1alpha1 "kusionstack.io/operating/apis/apps/v1alpha1"
+	utilspoddecoration "kusionstack.io/operating/pkg/controllers/utils/poddecoration"
+)
+
+func GetEffectiveDecorationsByCollaSet(
+	ctx context.Context,
+	c client.Client,
+	colla *appsv1alpha1.CollaSet,
+) (
+	podDecorations []*appsv1alpha1.PodDecoration, oldRevisions map[string]*appsv1alpha1.PodDecoration, err error) {
+
+	pdList := &appsv1alpha1.PodDecorationList{}
+	if err = c.List(ctx, pdList, &client.ListOptions{Namespace: colla.Namespace}); err != nil {
+		return
+	}
+	for i := range pdList.Items {
+		if isAffectedCollaSet(&pdList.Items[i], colla) {
+			podDecorations = append(podDecorations, &pdList.Items[i])
+		}
+	}
+	oldRevisions = map[string]*appsv1alpha1.PodDecoration{}
+	for _, pd := range podDecorations {
+		if pd.Status.CurrentRevision != "" && pd.Status.CurrentRevision != pd.Status.UpdatedRevision {
+			revision := &appsv1.ControllerRevision{}
+			if err = c.Get(ctx, types.NamespacedName{Namespace: colla.Namespace, Name: pd.Status.CurrentRevision}, revision); err != nil {
+				return nil, nil, fmt.Errorf("fail to get PodDecoration ControllerRevision %s/%s: %v", colla.Namespace, pd.Status.CurrentRevision, err)
+			}
+			oldPD, err := utilspoddecoration.GetPodDecorationFromRevision(revision)
+			if err != nil {
+				return nil, nil, err
+			}
+			oldRevisions[pd.Status.CurrentRevision] = oldPD
+		}
+	}
+	return
+}
+
+func isAffectedCollaSet(pd *appsv1alpha1.PodDecoration, colla *appsv1alpha1.CollaSet) bool {
+	if pd.Status.IsEffective == nil || !*pd.Status.IsEffective {
+		return false
+	}
+	sel, _ := metav1.LabelSelectorAsSelector(pd.Spec.Selector)
+	return sel.Matches(labels.Set(colla.Spec.Template.Labels))
+}

pkg/controllers/utils/poddecoration/lister.go → pkg/controllers/utils/poddecoration/getter.go

@@ -17,53 +17,13 @@ limitations under the License.
 package poddecoration
 
 import (
-	"context"
-	"fmt"
-	"sort"
-
-	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
-	"k8s.io/apimachinery/pkg/types"
-	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	appsv1alpha1 "kusionstack.io/operating/apis/apps/v1alpha1"
 )
 
-func GetEffectiveDecorationsByCollaSet(
-	ctx context.Context,
-	c client.Client,
-	colla *appsv1alpha1.CollaSet,
-) (
-	podDecorations []*appsv1alpha1.PodDecoration, oldRevisions map[string]*appsv1alpha1.PodDecoration, err error) {
-
-	pdList := &appsv1alpha1.PodDecorationList{}
-	if err = c.List(ctx, pdList, &client.ListOptions{Namespace: colla.Namespace}); err != nil {
-		return
-	}
-	for i := range pdList.Items {
-		if isAffectedCollaSet(&pdList.Items[i], colla) {
-			podDecorations = append(podDecorations, &pdList.Items[i])
-		}
-	}
-	oldRevisions = map[string]*appsv1alpha1.PodDecoration{}
-	for _, pd := range podDecorations {
-		if pd.Status.CurrentRevision != "" && pd.Status.CurrentRevision != pd.Status.UpdatedRevision {
-			revision := &appsv1.ControllerRevision{}
-			if err = c.Get(ctx, types.NamespacedName{Namespace: colla.Namespace, Name: pd.Status.CurrentRevision}, revision); err != nil {
-				return nil, nil, fmt.Errorf("fail to get PodDecoration ControllerRevision %s/%s: %v", colla.Namespace, pd.Status.CurrentRevision, err)
-			}
-			oldPD, err := GetPodDecorationFromRevision(revision)
-			if err != nil {
-				return nil, nil, err
-			}
-			oldRevisions[pd.Status.CurrentRevision] = oldPD
-		}
-	}
-	return
-}
-
 func GetPodEffectiveDecorations(pod *corev1.Pod, podDecorations []*appsv1alpha1.PodDecoration, oldRevisions map[string]*appsv1alpha1.PodDecoration) (res map[string]*appsv1alpha1.PodDecoration) {
 	type RevisionPD struct {
 		Revision string
@@ -84,9 +44,11 @@ func GetPodEffectiveDecorations(pod *corev1.Pod, podDecorations []*appsv1alpha1.
 			}
 			return
 		}
-		currentGroupPD[pd.Spec.InjectStrategy.Group] = &RevisionPD{
-			Revision: revision,
-			PD:       heaviestPD(current.PD, pd),
+		if lessPD(pd, current.PD) {
+			currentGroupPD[pd.Spec.InjectStrategy.Group] = &RevisionPD{
+				Revision: revision,
+				PD:       pd,
+			}
 		}
 	}
 	for i, pd := range podDecorations {
@@ -124,26 +86,3 @@ func GetPodEffectiveDecorations(pod *corev1.Pod, podDecorations []*appsv1alpha1.
 	}
 	return
 }
-
-func PickGroupTop(podDecorations []*appsv1alpha1.PodDecoration) (res []*appsv1alpha1.PodDecoration) {
-	sort.Sort(PodDecorations(podDecorations))
-	for i, pd := range podDecorations {
-		if i == 0 {
-			res = append(res, podDecorations[i])
-			continue
-		}
-		if pd.Spec.InjectStrategy.Group == res[len(res)-1].Spec.InjectStrategy.Group {
-			continue
-		}
-		res = append(res, podDecorations[i])
-	}
-	return
-}
-
-func isAffectedCollaSet(pd *appsv1alpha1.PodDecoration, colla *appsv1alpha1.CollaSet) bool {
-	if pd.Status.IsEffective == nil || !*pd.Status.IsEffective {
-		return false
-	}
-	sel, _ := metav1.LabelSelectorAsSelector(pd.Spec.Selector)
-	return sel.Matches(labels.Set(colla.Spec.Template.Labels))
-}

pkg/controllers/utils/poddecoration/patch_test.go

@@ -337,6 +337,106 @@ var _ = Describe("PodDecoration controller", func() {
 		Expect(pod.Spec.Tolerations[0].Key).Should(Equal("test"))
 		Expect(pod.Spec.Affinity).ShouldNot(BeNil())
 	})
+
+	It("test anno utils", func() {
+		pod := &v1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Annotations: map[string]string{
+					appsv1alpha1.AnnotationResourceDecorationRevision: "",
+				},
+				Labels: map[string]string{
+					"app": "foo",
+				},
+			},
+		}
+		i0Int32 := int32(0)
+		i1Int32 := int32(1)
+		pdA := &appsv1alpha1.PodDecoration{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pd-a",
+			},
+			Spec: appsv1alpha1.PodDecorationSpec{
+				InjectStrategy: appsv1alpha1.PodDecorationInjectStrategy{
+					Group:  "group-a",
+					Weight: &i0Int32,
+				},
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						"app": "foo",
+					},
+				},
+				Template: appsv1alpha1.PodDecorationPodTemplate{
+					Metadata: []*appsv1alpha1.PodDecorationPodTemplateMeta{
+						{
+							Labels: map[string]string{"inj": "group-a"},
+						},
+					},
+				},
+			},
+			Status: appsv1alpha1.PodDecorationStatus{
+				UpdatedRevision: "100",
+			},
+		}
+		pdB := &appsv1alpha1.PodDecoration{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pd-b",
+			},
+			Spec: appsv1alpha1.PodDecorationSpec{
+				InjectStrategy: appsv1alpha1.PodDecorationInjectStrategy{
+					Group:  "group-b",
+					Weight: &i1Int32,
+				},
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						"app": "foo",
+					},
+				},
+				Template: appsv1alpha1.PodDecorationPodTemplate{
+					Metadata: []*appsv1alpha1.PodDecorationPodTemplateMeta{
+						{
+							Labels: map[string]string{"inj": "group-b"},
+						},
+					},
+				},
+			},
+			Status: appsv1alpha1.PodDecorationStatus{
+				UpdatedRevision: "101",
+			},
+		}
+		pdC := &appsv1alpha1.PodDecoration{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pd-b",
+			},
+			Spec: appsv1alpha1.PodDecorationSpec{
+				InjectStrategy: appsv1alpha1.PodDecorationInjectStrategy{
+					Group:  "group-b",
+					Weight: &i1Int32,
+				},
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						"app": "foo",
+					},
+				},
+				Template: appsv1alpha1.PodDecorationPodTemplate{
+					Metadata: []*appsv1alpha1.PodDecorationPodTemplateMeta{
+						{
+							Labels: map[string]string{"inj": "group-b-new"},
+						},
+					},
+				},
+			},
+			Status: appsv1alpha1.PodDecorationStatus{
+				UpdatedRevision: "102",
+			},
+		}
+		pds := map[string]*appsv1alpha1.PodDecoration{
+			"100": pdA,
+			"101": pdB,
+		}
+		Expect(ShouldUpdateDecorationInfo(pod, pds)).Should(BeTrue())
+		Expect(PatchListOfDecorations(pod, pds)).Should(BeNil())
+		Expect(len(GetPodEffectiveDecorations(pod, []*appsv1alpha1.PodDecoration{pdA, pdC}, pds))).Should(Equal(2))
+	})
 })
 
 var _ = BeforeSuite(func() {