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

Signed-off-by: liheng.zms <liheng.zms@alibaba-inc.com>

apis/policy/v1alpha1/podunavailablebudget_types.go

@@ -35,9 +35,11 @@ const (
 	PubUpdateOperation PubOperation = "UPDATE"
 	PubDeleteOperation PubOperation = "DELETE"
 	PubEvictOperation  PubOperation = "EVICT"
-	// PubProtectTotalReplicas indicates the pub protected total replicas, rather than workload.spec.replicas.
+	// PubProtectTotalReplicasAnnotation is the target replicas.
-	// and must be used with pub.spec.selector.
+	// By default, PUB will get the target replicas through workload.spec.replicas. but there are some scenarios that may workload doesn't
-	PubProtectTotalReplicas = "pub.kruise.io/protect-total-replicas"
+	// implement scale subresources or Pod doesn't have workload management. In this scenario, you can set pub.kruise.io/protect-total-replicas
+	// in pub annotations to get the target replicas to realize the same effect of protection ability.
+	PubProtectTotalReplicasAnnotation = "pub.kruise.io/protect-total-replicas"
 	// Marked the pod will not be pub-protected, solving the scenario of force pod deletion
 	PodPubNoProtectionAnnotation = "pub.kruise.io/no-protect"
 )

pkg/control/pubcontrol/pub_control.go

@@ -78,6 +78,10 @@ func (c *commonControl) GetPodsForPub(pub *policyv1alpha1.PodUnavailableBudget)
 	if pub.Spec.TargetReference != nil {
 		ref := pub.Spec.TargetReference
 		matchedPods, expectedCount, err := c.controllerFinder.GetPodsForRef(ref.APIVersion, ref.Kind, pub.Namespace, ref.Name, true)
+		if value, _ := pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation]; value != "" {
+			count, _ := strconv.ParseInt(value, 10, 32)
+			expectedCount = int32(count)
+		}
 		return matchedPods, expectedCount, err
 	} else if pub.Spec.Selector == nil {
 		klog.InfoS("Pub spec.Selector could not be empty", "pub", klog.KObj(pub))
@@ -101,14 +105,14 @@ func (c *commonControl) GetPodsForPub(pub *policyv1alpha1.PodUnavailableBudget)
 			matchedPods = append(matchedPods, pod)
 		}
 	}
+	if value, _ := pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation]; value != "" {
+		expectedCount, _ := strconv.ParseInt(value, 10, 32)
+		return matchedPods, int32(expectedCount), nil
+	}
 	expectedCount, err := c.controllerFinder.GetExpectedScaleForPods(matchedPods)
 	if err != nil {
 		return nil, 0, err
 	}
-	if expectedCount == 0 && pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] != "" {
-		expectedCount, _ := strconv.ParseInt(pub.Annotations[policyv1alpha1.PubProtectTotalReplicas], 10, 32)
-		return matchedPods, int32(expectedCount), nil
-	}
 	return matchedPods, expectedCount, nil
 }
 

pkg/controller/podunavailablebudget/pub_controller_test.go

@@ -246,6 +246,67 @@ func TestPubReconcile(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "select matched pub.annotations[pub.kruise.io/protect-total-replicas]=15 and selector, 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.PubProtectTotalReplicasAnnotation] = "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 {
@@ -466,6 +527,56 @@ func TestPubReconcile(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "select matched deployment(replicas=1,maxSurge=0,maxUnavailable=30%), pub.kruise.io/protect-total-replicas=15 and pub(targetRef,maxUnavailable=30%)",
+			getPods: func(rs ...*apps.ReplicaSet) []*corev1.Pod {
+				var matchedPods []*corev1.Pod
+				for i := 0; int32(i) < 10; 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)
+				}
+				return matchedPods
+			},
+			getDeployment: func() *apps.Deployment {
+				obj := deploymentDemo.DeepCopy()
+				obj.Spec.Replicas = utilpointer.Int32(1)
+				return obj
+			},
+			getReplicaSet: func() []*apps.ReplicaSet {
+				obj1 := replicaSetDemo.DeepCopy()
+				obj1.Name = "nginx-rs-1"
+				return []*apps.ReplicaSet{obj1}
+			},
+			getPub: func() *policyv1alpha1.PodUnavailableBudget {
+				pub := pubDemo.DeepCopy()
+				pub.Spec.Selector = nil
+				pub.Spec.TargetReference = &policyv1alpha1.TargetReference{
+					APIVersion: "apps/v1",
+					Kind:       "Deployment",
+					Name:       "nginx",
+				}
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation] = "15"
+				return pub
+			},
+			expectPubStatus: func() policyv1alpha1.PodUnavailableBudgetStatus {
+				return policyv1alpha1.PodUnavailableBudgetStatus{
+					UnavailableAllowed: 0,
+					CurrentAvailable:   10,
+					DesiredAvailable:   10,
+					TotalReplicas:      15,
+				}
+			},
+		},
 		{
 			name: "select matched deployment(replicas=1,maxSurge=0,maxUnavailable=30%), and pub(targetRef,maxUnavailable=30%)",
 			getPods: func(rs ...*apps.ReplicaSet) []*corev1.Pod {
@@ -1114,7 +1225,7 @@ func TestPubReconcile(t *testing.T) {
 			getPub: func() *policyv1alpha1.PodUnavailableBudget {
 				pub := pubDemo.DeepCopy()
 
-				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "50"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation] = "50"
 				for i := 0; i < 10; i++ {
 					if i >= 0 && i < 5 {
 						pub.Status.UnavailablePods[fmt.Sprintf("test-pod-%d", i)] = metav1.Time{Time: time.Now().Add(-10 * time.Second)}

pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go

@@ -94,9 +94,9 @@ func (h *PodUnavailableBudgetCreateUpdateHandler) validatingPodUnavailableBudget
 			}
 		}
 	}
-	if replicasValue, ok := obj.Annotations[policyv1alpha1.PubProtectTotalReplicas]; ok {
+	if replicasValue, ok := obj.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation]; 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)))
+			allErrs = append(allErrs, field.InternalError(field.NewPath("metadata"), fmt.Errorf("annotation[%s] is invalid", policyv1alpha1.PubProtectTotalReplicasAnnotation)))
 		}
 	}
 

pkg/webhook/podunavailablebudget/validating/pub_validating_test.go

@@ -20,13 +20,11 @@ import (
 	"context"
 	"testing"
 
-	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
-
 	policyv1alpha1 "github.com/openkruise/kruise/apis/policy/v1alpha1"
-
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/intstr"
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 )
@@ -166,7 +164,7 @@ func TestValidatingPub(t *testing.T) {
 				pub := pubDemo.DeepCopy()
 				pub.Spec.Selector = nil
 				pub.Spec.MinAvailable = nil
-				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "%%"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation] = "%%"
 				return pub
 			},
 			expectErrList: 1,
@@ -177,7 +175,7 @@ func TestValidatingPub(t *testing.T) {
 				pub := pubDemo.DeepCopy()
 				pub.Spec.Selector = nil
 				pub.Spec.MinAvailable = nil
-				pub.Annotations[policyv1alpha1.PubProtectTotalReplicas] = "1000"
+				pub.Annotations[policyv1alpha1.PubProtectTotalReplicasAnnotation] = "1000"
 				return pub
 			},
 			expectErrList: 0,