refactor, PodIDAndTypesMap performance optimization

KusionStack · Aug 17, 2023 · fd63316 · fd63316
1 parent bc3217b
commit fd63316
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 103 deletions.
diff --git a/apis/apps/v1alpha1/well_known_labels.go b/apis/apps/v1alpha1/well_known_labels.go
@@ -18,6 +18,8 @@ package v1alpha1
 
 // pod ops lifecyle labels
 const (
+	ControlledByPodOpsLifecycle = "podopslifecycle.kusionstack.io/control" // indicate a pod is controlled by podopslifecycle
+
 	PodOperatingLabelPrefix           = "operating.podopslifecycle.kusionstack.io"            // indicate a pod is operating
 	PodOperationTypeLabelPrefix       = "operation-type.podopslifecycle.kusionstack.io"       // indicate the type of operation
 	PodOperationPermissionLabelPrefix = "operation-permission.podopslifecycle.kusionstack.io" // indicate the permission of operation
@@ -44,6 +46,6 @@ const (
 
 var (
 	WellKnownLabelPrefixesWithID = []string{PodOperatingLabelPrefix, PodOperationTypeLabelPrefix, PodPreCheckLabelPrefix, PodPreCheckedLabelPrefix,
-		PodPrepareLabelPrefix, PodUndoOperationTypeLabelPrefix, PodOperateLabelPrefix, PodOperatedLabelPrefix, PodPostCheckLabelPrefix,
+		PodPrepareLabelPrefix, PodDoneOperationTypeLabelPrefix, PodUndoOperationTypeLabelPrefix, PodOperateLabelPrefix, PodOperatedLabelPrefix, PodPostCheckLabelPrefix,
 		PodPostCheckedLabelPrefix, PodCompleteLabelPrefix}
 )
diff --git a/pkg/controllers/podopslifecycle/podopslifecycle_controller_test.go b/pkg/controllers/podopslifecycle/podopslifecycle_controller_test.go
@@ -166,6 +166,7 @@ var _ = Describe("podopslifecycle controller", func() {
 		}}
 
 		pod.ObjectMeta.Labels = map[string]string{
+			fmt.Sprintf("%s/%s", v1alpha1.PodOperateLabelPrefix, id):       time,
 			fmt.Sprintf("%s/%s", v1alpha1.PodPreCheckLabelPrefix, id):      time,
 			fmt.Sprintf("%s/%s", v1alpha1.PodOperationTypeLabelPrefix, id): operationType,
 		}
@@ -191,7 +192,8 @@ var _ = Describe("podopslifecycle controller", func() {
 				Name:      "test",
 				Namespace: "default",
 				Labels: map[string]string{
-					fmt.Sprintf("%s/%s", v1alpha1.PodPrepareLabelPrefix, id): time,
+					fmt.Sprintf("%s/%s", v1alpha1.PodOperatingLabelPrefix, id): time,
+					fmt.Sprintf("%s/%s", v1alpha1.PodPrepareLabelPrefix, id):   time,
 				},
 			},
 			Spec: podSpec,
@@ -218,6 +220,7 @@ var _ = Describe("podopslifecycle controller", func() {
 				Name:      "test",
 				Namespace: "default",
 				Labels: map[string]string{
+					fmt.Sprintf("%s/%s", v1alpha1.PodOperateLabelPrefix, id):  time,
 					fmt.Sprintf("%s/%s", v1alpha1.PodCompleteLabelPrefix, id): time,
 				},
 			},
@@ -261,6 +264,7 @@ var _ = Describe("podopslifecycle controller", func() {
 		Expect(pod.Status.Conditions).To(HaveLen(0))
 
 		pod.ObjectMeta.Labels = map[string]string{
+			fmt.Sprintf("%s/%s", v1alpha1.PodOperateLabelPrefix, id):  time,
 			fmt.Sprintf("%s/%s", v1alpha1.PodCompleteLabelPrefix, id): time,
 		}
 		err = mgr.GetClient().Update(context.Background(), pod)
@@ -315,7 +319,7 @@ func (rsm *mockRuleSetManager) GetState(client.Client, client.Object) (checker.C
 	return *rsm.CheckState, nil
 }
 
-func TestPodOpsLifecycleController(t *testing.T) {
+func TestControlledByPodOpsLifecycleler(t *testing.T) {
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "podopslifecycle controller suite test")
 }
diff --git a/pkg/controllers/podopslifecycle/utils.go b/pkg/controllers/podopslifecycle/utils.go
@@ -21,6 +21,7 @@ import (
 	"strings"
 
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	"kusionstack.io/kafed/apis/apps/v1alpha1"
 )
@@ -30,34 +31,40 @@ func PodIDAndTypesMap(pod *v1.Pod) (map[string]map[string]string, map[string]int
 	idToLabelsMap := map[string]map[string]string{}
 	typeToNumsMap := map[string]int{}
 
-	for k, v := range pod.Labels {
-		if strings.HasPrefix(k, v1alpha1.PodOperationTypeLabelPrefix) {
-			if _, ok := typeToNumsMap[v]; !ok {
-				typeToNumsMap[v] = 1
-			} else {
-				typeToNumsMap[v] = typeToNumsMap[v] + 1
+	ids := sets.String{}
+	for k := range pod.Labels {
+		if strings.HasPrefix(k, v1alpha1.PodOperatingLabelPrefix) || strings.HasPrefix(k, v1alpha1.PodOperateLabelPrefix) {
+			s := strings.Split(k, "/")
+			if len(s) < 2 {
+				return nil, nil, fmt.Errorf("invalid label %s", k)
 			}
+			ids.Insert(s[1])
 		}
+	}
 
-		for _, label := range v1alpha1.WellKnownLabelPrefixesWithID {
-			if !strings.HasPrefix(k, label) {
-				continue
+	for id := range ids {
+		operationType, ok := pod.Labels[fmt.Sprintf("%s/%s", v1alpha1.PodOperationTypeLabelPrefix, id)]
+		if ok {
+			if _, ok := typeToNumsMap[operationType]; !ok {
+				typeToNumsMap[operationType] = 1
+			} else {
+				typeToNumsMap[operationType] = typeToNumsMap[operationType] + 1
 			}
+		}
 
-			s := strings.Split(k, "/")
-			if len(s) < 2 {
-				return nil, nil, fmt.Errorf("invalid label %s", k)
+		for _, prefix := range v1alpha1.WellKnownLabelPrefixesWithID {
+			label := fmt.Sprintf("%s/%s", prefix, id)
+			value, ok := pod.Labels[label]
+			if !ok {
+				continue
 			}
-			id := s[1]
 
 			labelsMap, ok := idToLabelsMap[id]
 			if !ok {
 				labelsMap = make(map[string]string)
 				idToLabelsMap[id] = labelsMap
 			}
-			labelsMap[label] = v
-
-			break
+			labelsMap[prefix] = value
 		}
 	}
 	return idToLabelsMap, typeToNumsMap, nil

diff --git a/pkg/utils/common.go b/pkg/utils/common.go
@@ -26,6 +26,6 @@ func ControlledByKafed(obj client.Object) bool {
 	if obj == nil || obj.GetLabels() == nil {
 		return false
 	}
-	v, ok := obj.GetLabels()[v1alpha1.KafedSystemLabel]
+	v, ok := obj.GetLabels()[v1alpha1.ControlledByPodOpsLifecycle]
 	return ok && v == "true"
 }
diff --git a/pkg/webhook/server/generic/pod/opslifecycle/mutating.go b/pkg/webhook/server/generic/pod/opslifecycle/mutating.go
@@ -80,7 +80,6 @@ func (lc *OpsLifecycle) Mutating(ctx context.Context, c client.Client, oldPod, n
 
 		if _, ok := labels[v1alpha1.PodOperateLabelPrefix]; ok {
 			operateCount++
-			continue
 		}
 
 		if _, ok := labels[v1alpha1.PodOperatedLabelPrefix]; ok { // operated
@@ -106,7 +105,23 @@ func (lc *OpsLifecycle) Mutating(ctx context.Context, c client.Client, oldPod, n
 		}
 	}
 
-	if operatingCount != 0 { // wait for all operatings to be done
+	if operatingCount != 0 { // wait for all operations to be done
+		return nil
+	}
+
+	if completeCount == len(newIdToLabelsMap) { // all operations are done
+		satisfied, _, err := lc.satisfyExpectedFinalizers(newPod)
+		if err != nil {
+			return err
+		}
+		if satisfied { // all operations are done and all expected finalizers are satisfied, then remove all unuseful labels, and add service available label
+			for id := range newIdToLabelsMap {
+				for _, v := range []string{v1alpha1.PodOperateLabelPrefix, v1alpha1.PodOperatedLabelPrefix, v1alpha1.PodDoneOperationTypeLabelPrefix, v1alpha1.PodPostCheckLabelPrefix, v1alpha1.PodPostCheckedLabelPrefix, v1alpha1.PodCompleteLabelPrefix} {
+					delete(newPod.Labels, fmt.Sprintf("%s/%s", v, id))
+				}
+			}
+			lc.addLabelWithTime(newPod, v1alpha1.PodServiceAvailableLabel)
+		}
 		return nil
 	}
 
@@ -116,36 +131,23 @@ func (lc *OpsLifecycle) Mutating(ctx context.Context, c client.Client, oldPod, n
 			return err
 		}
 
-		for id := range newIdToLabelsMap {
-			for _, v := range []string{v1alpha1.PodPreCheckLabelPrefix, v1alpha1.PodPreCheckedLabelPrefix, v1alpha1.PodPrepareLabelPrefix, v1alpha1.PodOperateLabelPrefix} {
+		for id, labels := range newIdToLabelsMap {
+			for _, v := range []string{v1alpha1.PodPreCheckLabelPrefix, v1alpha1.PodPreCheckedLabelPrefix, v1alpha1.PodPrepareLabelPrefix} {
 				delete(newPod.Labels, fmt.Sprintf("%s/%s", v, id))
 			}
-			lc.addLabelWithTime(newPod, fmt.Sprintf("%s/%s", v1alpha1.PodOperatedLabelPrefix, id))
+
+			if _, ok := labels[v1alpha1.PodOperatedLabelPrefix]; !ok {
+				lc.addLabelWithTime(newPod, fmt.Sprintf("%s/%s", v1alpha1.PodOperatedLabelPrefix, id))
+			}
 
 			t, ok := oldIdToLabelsMap[id][v1alpha1.PodOperationTypeLabelPrefix]
 			if !ok {
-				return fmt.Errorf("pod %s/%s label %s not found", oldPod.Namespace, oldPod.Name, fmt.Sprintf("%s/%s", v1alpha1.PodPreCheckedLabelPrefix, id))
+				continue
 			}
-
 			delete(newPod.Labels, fmt.Sprintf("%s/%s", v1alpha1.PodOperationPermissionLabelPrefix, t))
 			newPod.Labels[fmt.Sprintf("%s/%s", v1alpha1.PodDoneOperationTypeLabelPrefix, id)] = t
 		}
 	}
 
-	if completeCount == len(newIdToLabelsMap) { // all operations are done
-		satisfied, _, err := lc.satisfyExpectedFinalizers(newPod)
-		if err != nil {
-			return err
-		}
-		if satisfied { // all operations are done and all expected finalizers are satisfied, then remove all unuseful labels, and add service available label
-			for id := range newIdToLabelsMap {
-				for _, v := range []string{v1alpha1.PodOperatedLabelPrefix, v1alpha1.PodDoneOperationTypeLabelPrefix, v1alpha1.PodPostCheckLabelPrefix, v1alpha1.PodPostCheckedLabelPrefix, v1alpha1.PodCompleteLabelPrefix} {
-					delete(newPod.Labels, fmt.Sprintf("%s/%s", v, id))
-				}
-			}
-			lc.addLabelWithTime(newPod, v1alpha1.PodServiceAvailableLabel)
-		}
-	}
-
 	return nil
 }