support pub pub.kruise.io/disable-fetch-replicas-from-workload=true

Signed-off-by: liheng.zms <liheng.zms@alibaba-inc.com>
openkruise · Sep 24, 2024 · 4a1339e · 4a1339e
1 parent 4f04e93
commit 4a1339e
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 1 deletion.
diff --git a/apis/policy/v1alpha1/podunavailablebudget_types.go b/apis/policy/v1alpha1/podunavailablebudget_types.go
@@ -37,7 +37,10 @@ const (
 	PubEvictOperation  PubOperation = "EVICT"
 	// PubProtectTotalReplicas indicates the pub protected total replicas, rather than workload.spec.replicas.
 	// and must be used with pub.spec.selector.
-	PubProtectTotalReplicas = "pub.kruise.io/protect-total-replicas"
+	// Even if pub.kruise.io/protect-total-replicas is configured, pub will still prioritize getting total replicas through the pod ownerRef workload.
+	// If you only want to use pub.kruise.io/protect-total-replicas, please set pub.kruise.io/disable-fetch-replicas-from-workload=true.
+	PubProtectTotalReplicas             = "pub.kruise.io/protect-total-replicas"
+	PubDisableFetchReplicasFromWorkload = "pub.kruise.io/disable-fetch-replicas-from-workload"
 	// Marked the pod will not be pub-protected, solving the scenario of force pod deletion
 	PodPubNoProtectionAnnotation = "pub.kruise.io/no-protect"
 )

diff --git a/pkg/control/pubcontrol/pub_control.go b/pkg/control/pubcontrol/pub_control.go
@@ -101,6 +101,10 @@ func (c *commonControl) GetPodsForPub(pub *policyv1alpha1.PodUnavailableBudget)
 			matchedPods = append(matchedPods, pod)
 		}
 	}
+	if v, _ := pub.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload]; v == "true" {
+		expectedCount, _ := strconv.ParseInt(pub.Annotations[policyv1alpha1.PubProtectTotalReplicas], 10, 32)
+		return matchedPods, int32(expectedCount), nil
+	}
 	expectedCount, err := c.controllerFinder.GetExpectedScaleForPods(matchedPods)
 	if err != nil {
 		return nil, 0, err

diff --git a/pkg/controller/podunavailablebudget/pub_controller_test.go b/pkg/controller/podunavailablebudget/pub_controller_test.go
@@ -246,6 +246,68 @@ func TestPubReconcile(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "select matched pub.annotations[pub.kruise.io/protect-total-replicas]=15, selector and maxUnavailable 30%",
+			getPods: func(rs ...*apps.ReplicaSet) []*corev1.Pod {
+				var matchedPods []*corev1.Pod
+				for i := 0; int32(i) < 5; i++ {
+					pod := podDemo.DeepCopy()
+					pod.OwnerReferences = []metav1.OwnerReference{
+						{
+							APIVersion: "apps/v1",
+							Kind:       "ReplicaSet",
+							Name:       rs[0].Name,
+							UID:        rs[0].UID,
+							Controller: utilpointer.BoolPtr(true),
+						},
+					}
+					pod.Name = fmt.Sprintf("%s-%d", pod.Name, i)
+					matchedPods = append(matchedPods, pod)
+				}
+				for i := 5; int32(i) < 10; i++ {
+					pod := podDemo.DeepCopy()
+					pod.OwnerReferences = []metav1.OwnerReference{
+						{
+							APIVersion: "apps/v1",
+							Kind:       "ReplicaSet",
+							Name:       rs[1].Name,
+							UID:        rs[1].UID,
+							Controller: utilpointer.BoolPtr(true),
+						},
+					}
+					pod.Name = fmt.Sprintf("%s-%d", pod.Name, i)
+					matchedPods = append(matchedPods, pod)
+				}
+				return matchedPods
+			},
+			getDeployment: func() *apps.Deployment {
+				obj := deploymentDemo.DeepCopy()
+				obj.Spec.Replicas = utilpointer.Int32(0)
+				return obj
+			},
+			getReplicaSet: func() []*apps.ReplicaSet {
+				obj1 := replicaSetDemo.DeepCopy()
+				obj1.Name = "nginx-rs-1"
+				obj2 := replicaSetDemo.DeepCopy()
+				obj2.Name = "nginx-rs-2"
+				obj2.UID = "a34b0453-3426-4685-a79c-752e7062a523"
+				return []*apps.ReplicaSet{obj1, obj2}
+			},
+			getPub: func() *policyv1alpha1.PodUnavailableBudget {
+				pub := pubDemo.DeepCopy()
+				pub.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload] = "true"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "15"
+				return pub
+			},
+			expectPubStatus: func() policyv1alpha1.PodUnavailableBudgetStatus {
+				return policyv1alpha1.PodUnavailableBudgetStatus{
+					UnavailableAllowed: 0,
+					CurrentAvailable:   10,
+					DesiredAvailable:   10,
+					TotalReplicas:      15,
+				}
+			},
+		},
 		{
 			name: "select matched deployment(replicas=10,maxSurge=30%,maxUnavailable=0), and pub(selector,maxUnavailable=30%)",
 			getPods: func(rs ...*apps.ReplicaSet) []*corev1.Pod {

diff --git a/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go b/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go
@@ -94,6 +94,16 @@ func (h *PodUnavailableBudgetCreateUpdateHandler) validatingPodUnavailableBudget
 			}
 		}
 	}
+	if v, _ := obj.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload]; v == "true" {
+		if _, ok := obj.Annotations[policyv1alpha1.PubProtectTotalReplicas]; !ok {
+			allErrs = append(allErrs, field.InternalError(field.NewPath("metadata"), fmt.Errorf("when annotation[%s]='true', annotation[%s] must be set",
+				policyv1alpha1.PubDisableFetchReplicasFromWorkload, policyv1alpha1.PubProtectTotalReplicas)))
+		}
+		if obj.Spec.TargetReference != nil {
+			allErrs = append(allErrs, field.InternalError(field.NewPath("metadata"), fmt.Errorf("when annotation[%s]='true', only spec.selector can be configured",
+				policyv1alpha1.PubDisableFetchReplicasFromWorkload)))
+		}
+	}
 	if replicasValue, ok := obj.Annotations[policyv1alpha1.PubProtectTotalReplicas]; ok {
 		if _, err := strconv.ParseInt(replicasValue, 10, 32); err != nil {
 			allErrs = append(allErrs, field.InternalError(field.NewPath("metadata"), fmt.Errorf("annotation[%s] is invalid", policyv1alpha1.PubProtectTotalReplicas)))

diff --git a/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go b/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go
@@ -182,6 +182,41 @@ func TestValidatingPub(t *testing.T) {
 			},
 			expectErrList: 0,
 		},
+		{
+			name: "invalid pub annotation[pub.kruise.io/disable-fetch-replicas-from-workload]='true', but pub.kruise.io/protect-total-replicas not set",
+			pub: func() *policyv1alpha1.PodUnavailableBudget {
+				pub := pubDemo.DeepCopy()
+				pub.Spec.TargetReference = nil
+				pub.Spec.MinAvailable = nil
+				pub.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload] = "true"
+				return pub
+			},
+			expectErrList: 1,
+		},
+		{
+			name: "invalid pub annotation[pub.kruise.io/disable-fetch-replicas-from-workload]='true', but set spec.TargetReference",
+			pub: func() *policyv1alpha1.PodUnavailableBudget {
+				pub := pubDemo.DeepCopy()
+				pub.Spec.Selector = nil
+				pub.Spec.MinAvailable = nil
+				pub.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload] = "true"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "1000"
+				return pub
+			},
+			expectErrList: 1,
+		},
+		{
+			name: "valid pub annotation[pub.kruise.io/disable-fetch-replicas-from-workload]='true'",
+			pub: func() *policyv1alpha1.PodUnavailableBudget {
+				pub := pubDemo.DeepCopy()
+				pub.Spec.TargetReference = nil
+				pub.Spec.MinAvailable = nil
+				pub.Annotations[policyv1alpha1.PubDisableFetchReplicasFromWorkload] = "true"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "1000"
+				return pub
+			},
+			expectErrList: 0,
+		},
 	}
 
 	decoder := admission.NewDecoder(scheme)