diff --git a/apis/policy/v1alpha1/podunavailablebudget_types.go b/apis/policy/v1alpha1/podunavailablebudget_types.go index 4d07eb3fa0..54383203f2 100644 --- a/apis/policy/v1alpha1/podunavailablebudget_types.go +++ b/apis/policy/v1alpha1/podunavailablebudget_types.go @@ -24,6 +24,14 @@ import ( // EDIT THIS FILE! THIS IS SCAFFOLDING FOR YOU TO OWN! // NOTE: json tags are required. Any new fields you add must have json tags for the fields to be serialized. +const ( + // PubProtectOperationAnnotation indicates the pub protected Operation[DELETE,UPDATE] + // the following indicates the pub only protect DELETE,UPDATE Operation + // annotations[kruise.io/pub-protect-operations]=DELETE,UPDATE + // if the annotations do not exist, the default DELETE and UPDATE are protected + PubProtectOperationAnnotation = "kruise.io/pub-protect-operations" +) + // PodUnavailableBudgetSpec defines the desired state of PodUnavailableBudget type PodUnavailableBudgetSpec struct { // Selector label query over pods managed by the budget diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index 5baafbb8de..a6c4eac58b 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -42,6 +42,14 @@ rules: - '*' verbs: - list +- apiGroups: + - '*' + resources: + - '*/scale' + verbs: + - get + - list + - watch - apiGroups: - admissionregistration.k8s.io resources: diff --git a/main.go b/main.go index 4cdda0c01f..549e611719 100644 --- a/main.go +++ b/main.go @@ -24,6 +24,7 @@ import ( "os" "time" + "github.com/openkruise/kruise/pkg/util/controllerfinder" "github.com/spf13/pflag" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" @@ -149,6 +150,11 @@ func main() { setupLog.Error(err, "unable to start manager") os.Exit(1) } + err = controllerfinder.InitControllerFinder(mgr) + if err != nil { + setupLog.Error(err, "unable to start ControllerFinder") + os.Exit(1) + } setupLog.Info("register field index") if err := fieldindex.RegisterFieldIndexes(mgr.GetCache()); err != nil { diff --git a/pkg/control/pubcontrol/api.go b/pkg/control/pubcontrol/api.go index 33c67de3f0..d8019b473e 100644 --- a/pkg/control/pubcontrol/api.go +++ b/pkg/control/pubcontrol/api.go @@ -44,6 +44,6 @@ type PubControl interface { } func NewPubControl(client client.Client) PubControl { - controllerFinder := controllerfinder.NewControllerFinder(client) + controllerFinder := controllerfinder.Finder return &commonControl{controllerFinder: controllerFinder, Client: client} } diff --git a/pkg/control/pubcontrol/utils.go b/pkg/control/pubcontrol/utils.go index 8b5f50b172..0e5bd3a97b 100644 --- a/pkg/control/pubcontrol/utils.go +++ b/pkg/control/pubcontrol/utils.go @@ -24,12 +24,9 @@ import ( policyv1alpha1 "github.com/openkruise/kruise/apis/policy/v1alpha1" kubeClient "github.com/openkruise/kruise/pkg/client" "github.com/openkruise/kruise/pkg/util" - utilclient "github.com/openkruise/kruise/pkg/util/client" - "github.com/openkruise/kruise/pkg/util/controllerfinder" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" @@ -55,7 +52,7 @@ type Operation string const ( UpdateOperation = "UPDATE" - //DeleteOperation = "DELETE" + DeleteOperation = "DELETE" // Marked pods will not be pub-protected, solving the scenario of force pod deletion PodPubNoProtectionAnnotation = "pub.kruise.io/no-protect" @@ -203,40 +200,6 @@ func isPodRecordedInPub(podName string, pub *policyv1alpha1.PodUnavailableBudget return false } -func GetPubForWorkload(c client.Client, workload *controllerfinder.ScaleAndSelector) (*policyv1alpha1.PodUnavailableBudget, error) { - pubList := &policyv1alpha1.PodUnavailableBudgetList{} - if err := c.List(context.TODO(), pubList, &client.ListOptions{Namespace: workload.Metadata.Namespace}, utilclient.DisableDeepCopy); err != nil { - return nil, err - } - for i := range pubList.Items { - pub := &pubList.Items[i] - // if targetReference isn't nil, priority to take effect - if pub.Spec.TargetReference != nil { - // belongs the same workload - if IsReferenceEqual(&policyv1alpha1.TargetReference{ - APIVersion: workload.APIVersion, - Kind: workload.Kind, - Name: workload.Name, - }, pub.Spec.TargetReference) { - return pub, nil - } - } else { - // This error is irreversible, so continue - labelSelector, err := util.GetFastLabelSelector(pub.Spec.Selector) - if err != nil { - continue - } - // If a PUB with a nil or empty selector creeps in, it should match nothing, not everything. - if labelSelector.Empty() || !labelSelector.Matches(labels.Set(workload.TempLabels)) { - continue - } - return pub, nil - } - } - klog.V(6).Infof("could not find PodUnavailableBudget for workload %s in namespace %s with labels: %v", workload.Name, workload.Metadata.Namespace, workload.TempLabels) - return nil, nil -} - // check APIVersion, Kind, Name func IsReferenceEqual(ref1, ref2 *policyv1alpha1.TargetReference) bool { gv1, err := schema.ParseGroupVersion(ref1.APIVersion) diff --git a/pkg/controller/cloneset/sync/api.go b/pkg/controller/cloneset/sync/api.go index 2a9e683529..94209347e4 100644 --- a/pkg/controller/cloneset/sync/api.go +++ b/pkg/controller/cloneset/sync/api.go @@ -58,7 +58,7 @@ func New(c client.Client, recorder record.EventRecorder) Interface { inplaceControl: inplaceupdate.New(c, clonesetutils.RevisionAdapterImpl), lifecycleControl: lifecycle.New(c), recorder: recorder, - controllerFinder: controllerfinder.NewControllerFinder(c), + controllerFinder: controllerfinder.Finder, pubControl: pubcontrol.NewPubControl(c), } } diff --git a/pkg/controller/cloneset/sync/cloneset_update_test.go b/pkg/controller/cloneset/sync/cloneset_update_test.go index f65ae72069..7f1dfe352a 100644 --- a/pkg/controller/cloneset/sync/cloneset_update_test.go +++ b/pkg/controller/cloneset/sync/cloneset_update_test.go @@ -588,7 +588,7 @@ func TestUpdate(t *testing.T) { lifecycle.New(fakeClient), inplaceupdate.New(fakeClient, clonesetutils.RevisionAdapterImpl), record.NewFakeRecorder(10), - controllerfinder.NewControllerFinder(fakeClient), + &controllerfinder.ControllerFinder{Client: fakeClient}, pubcontrol.NewPubControl(fakeClient), } currentRevision := mc.updateRevision diff --git a/pkg/controller/persistentpodstate/persistent_pod_state_controller.go b/pkg/controller/persistentpodstate/persistent_pod_state_controller.go index bdc51c9893..100650fa39 100644 --- a/pkg/controller/persistentpodstate/persistent_pod_state_controller.go +++ b/pkg/controller/persistentpodstate/persistent_pod_state_controller.go @@ -79,7 +79,7 @@ func newReconciler(mgr manager.Manager) reconcile.Reconciler { return &ReconcilePersistentPodState{ Client: cli, scheme: mgr.GetScheme(), - finder: controllerfinder.NewControllerFinder(cli), + finder: controllerfinder.Finder, } } diff --git a/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go b/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go index 7131376a0d..48a2e65cca 100644 --- a/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go +++ b/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go @@ -519,7 +519,7 @@ func TestReconcilePersistentPodState(t *testing.T) { fakeClient := clientBuilder.Build() reconciler := ReconcilePersistentPodState{ Client: fakeClient, - finder: controllerfinder.NewControllerFinder(fakeClient), + finder: &controllerfinder.ControllerFinder{Client: fakeClient}, } if _, err := reconciler.Reconcile(context.TODO(), request); err != nil { t.Fatalf("reconcile failed, err: %v", err) diff --git a/pkg/controller/podunavailablebudget/podunavailablebudget_controller.go b/pkg/controller/podunavailablebudget/podunavailablebudget_controller.go index 17523eda5b..564b49f64c 100644 --- a/pkg/controller/podunavailablebudget/podunavailablebudget_controller.go +++ b/pkg/controller/podunavailablebudget/podunavailablebudget_controller.go @@ -103,7 +103,7 @@ func newReconciler(mgr manager.Manager) reconcile.Reconciler { Client: mgr.GetClient(), Scheme: mgr.GetScheme(), recorder: mgr.GetEventRecorderFor("podunavailablebudget-controller"), - controllerFinder: controllerfinder.NewControllerFinder(mgr.GetClient()), + controllerFinder: controllerfinder.Finder, pubControl: pubcontrol.NewPubControl(mgr.GetClient()), } } @@ -208,6 +208,7 @@ type ReconcilePodUnavailableBudget struct { // +kubebuilder:rbac:groups=policy.kruise.io,resources=podunavailablebudgets,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=policy.kruise.io,resources=podunavailablebudgets/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=*,resources=*/scale,verbs=get;list;watch // pkg/controller/cloneset/cloneset_controller.go Watch for changes to CloneSet func (r *ReconcilePodUnavailableBudget) Reconcile(_ context.Context, req ctrl.Request) (ctrl.Result, error) { diff --git a/pkg/controller/podunavailablebudget/pub_controller_test.go b/pkg/controller/podunavailablebudget/pub_controller_test.go index d2bfad9d55..c9de06127d 100644 --- a/pkg/controller/podunavailablebudget/pub_controller_test.go +++ b/pkg/controller/podunavailablebudget/pub_controller_test.go @@ -767,10 +767,11 @@ func TestPubReconcile(t *testing.T) { t.Fatalf("create pod failed: %s", err.Error()) } } + controllerfinder.Finder = &controllerfinder.ControllerFinder{Client: fakeClient} reconciler := ReconcilePodUnavailableBudget{ Client: fakeClient, recorder: record.NewFakeRecorder(10), - controllerFinder: controllerfinder.NewControllerFinder(fakeClient), + controllerFinder: &controllerfinder.ControllerFinder{Client: fakeClient}, pubControl: pubcontrol.NewPubControl(fakeClient), } diff --git a/pkg/controller/podunavailablebudget/pub_pod_event_handler.go b/pkg/controller/podunavailablebudget/pub_pod_event_handler.go index d230811b2d..d25f173bc8 100644 --- a/pkg/controller/podunavailablebudget/pub_pod_event_handler.go +++ b/pkg/controller/podunavailablebudget/pub_pod_event_handler.go @@ -25,9 +25,11 @@ import ( policyv1alpha1 "github.com/openkruise/kruise/apis/policy/v1alpha1" "github.com/openkruise/kruise/pkg/control/pubcontrol" "github.com/openkruise/kruise/pkg/util" + utilclient "github.com/openkruise/kruise/pkg/util/client" "github.com/openkruise/kruise/pkg/util/controllerfinder" apps "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/runtime" "k8s.io/apimachinery/pkg/types" @@ -45,7 +47,7 @@ var _ handler.EventHandler = &enqueueRequestForPod{} func newEnqueueRequestForPod(c client.Client) handler.EventHandler { e := &enqueueRequestForPod{client: c} - e.controllerFinder = controllerfinder.NewControllerFinder(c) + e.controllerFinder = controllerfinder.Finder e.pubControl = pubcontrol.NewPubControl(c) return e } @@ -74,9 +76,12 @@ func (p *enqueueRequestForPod) addPod(q workqueue.RateLimitingInterface, obj run if !ok { return } - - // reconcile pub - pub, _ := p.pubControl.GetPubForPod(pod) + var pub *policyv1alpha1.PodUnavailableBudget + if pod.Annotations[pubcontrol.PodRelatedPubAnnotation] == "" { + pub, _ = GetPubForPod(p.client, pod) + } else { + pub, _ = p.pubControl.GetPubForPod(pod) + } if pub == nil { return } @@ -89,6 +94,49 @@ func (p *enqueueRequestForPod) addPod(q workqueue.RateLimitingInterface, obj run }) } +func GetPubForPod(c client.Client, pod *corev1.Pod) (*policyv1alpha1.PodUnavailableBudget, error) { + ref := metav1.GetControllerOf(pod) + if ref == nil { + return nil, nil + } + workload, err := controllerfinder.Finder.GetScaleAndSelectorForRef(ref.APIVersion, ref.Kind, pod.Namespace, ref.Name, "") + if err != nil { + return nil, err + } else if workload == nil { + return nil, nil + } + pubList := &policyv1alpha1.PodUnavailableBudgetList{} + if err = c.List(context.TODO(), pubList, &client.ListOptions{Namespace: pod.Namespace}, utilclient.DisableDeepCopy); err != nil { + return nil, err + } + for i := range pubList.Items { + pub := &pubList.Items[i] + // if targetReference isn't nil, priority to take effect + if pub.Spec.TargetReference != nil { + // belongs the same workload + if pubcontrol.IsReferenceEqual(&policyv1alpha1.TargetReference{ + APIVersion: workload.APIVersion, + Kind: workload.Kind, + Name: workload.Name, + }, pub.Spec.TargetReference) { + return pub, nil + } + } else { + // This error is irreversible, so continue + labelSelector, err := util.GetFastLabelSelector(pub.Spec.Selector) + if err != nil { + continue + } + // If a PUB with a nil or empty selector creeps in, it should match nothing, not everything. + if labelSelector.Empty() || !labelSelector.Matches(labels.Set(pod.Labels)) { + continue + } + return pub, nil + } + } + return nil, nil +} + func (p *enqueueRequestForPod) updatePod(q workqueue.RateLimitingInterface, old, cur runtime.Object) { newPod := cur.(*corev1.Pod) oldPod := old.(*corev1.Pod) diff --git a/pkg/controller/workloadspread/reschedule_test.go b/pkg/controller/workloadspread/reschedule_test.go index 69853d0576..a7cd888523 100644 --- a/pkg/controller/workloadspread/reschedule_test.go +++ b/pkg/controller/workloadspread/reschedule_test.go @@ -268,9 +268,11 @@ func TestRescheduleSubset(t *testing.T) { } reconciler := ReconcileWorkloadSpread{ - Client: fakeClient, - recorder: record.NewFakeRecorder(10), - controllerFinder: controllerfinder.NewControllerFinder(fakeClient), + Client: fakeClient, + recorder: record.NewFakeRecorder(10), + controllerFinder: &controllerfinder.ControllerFinder{ + Client: fakeClient, + }, } err := reconciler.syncWorkloadSpread(workloadSpread) diff --git a/pkg/controller/workloadspread/workloadspread_controller.go b/pkg/controller/workloadspread/workloadspread_controller.go index 9dc4e9a8bd..92236b433b 100644 --- a/pkg/controller/workloadspread/workloadspread_controller.go +++ b/pkg/controller/workloadspread/workloadspread_controller.go @@ -154,7 +154,7 @@ func newReconciler(mgr manager.Manager) reconcile.Reconciler { Client: cli, scheme: mgr.GetScheme(), recorder: mgr.GetEventRecorderFor(controllerName), - controllerFinder: controllerfinder.NewControllerFinder(mgr.GetClient()), + controllerFinder: controllerfinder.Finder, } } diff --git a/pkg/controller/workloadspread/workloadspread_controller_test.go b/pkg/controller/workloadspread/workloadspread_controller_test.go index 5cbc6a808c..71bec9a171 100644 --- a/pkg/controller/workloadspread/workloadspread_controller_test.go +++ b/pkg/controller/workloadspread/workloadspread_controller_test.go @@ -1409,7 +1409,7 @@ func TestWorkloadSpreadReconcile(t *testing.T) { reconciler := ReconcileWorkloadSpread{ Client: fakeClient, recorder: record.NewFakeRecorder(10), - controllerFinder: controllerfinder.NewControllerFinder(fakeClient), + controllerFinder: &controllerfinder.ControllerFinder{Client: fakeClient}, } err := reconciler.syncWorkloadSpread(workloadSpread) @@ -1609,7 +1609,7 @@ func TestDelayReconcile(t *testing.T) { reconciler := ReconcileWorkloadSpread{ Client: fakeClient, recorder: record.NewFakeRecorder(10), - controllerFinder: controllerfinder.NewControllerFinder(fakeClient), + controllerFinder: &controllerfinder.ControllerFinder{Client: fakeClient}, } durationStore = requeueduration.DurationStore{} @@ -2031,7 +2031,7 @@ func TestManagerExistingPods(t *testing.T) { reconciler := ReconcileWorkloadSpread{ Client: fakeClient, recorder: record.NewFakeRecorder(10), - controllerFinder: controllerfinder.NewControllerFinder(fakeClient), + controllerFinder: &controllerfinder.ControllerFinder{Client: fakeClient}, } durationStore = requeueduration.DurationStore{} diff --git a/pkg/util/controllerfinder/controller_finder.go b/pkg/util/controllerfinder/controller_finder.go index a561f0ef47..569c1a2261 100644 --- a/pkg/util/controllerfinder/controller_finder.go +++ b/pkg/util/controllerfinder/controller_finder.go @@ -24,12 +24,38 @@ import ( apps "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/discovery" + "k8s.io/client-go/dynamic" + clientset "k8s.io/client-go/kubernetes" + scaleclient "k8s.io/client-go/scale" "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/manager" ) +var Finder *ControllerFinder + +func InitControllerFinder(mgr manager.Manager) error { + Finder = &ControllerFinder{ + Client: mgr.GetClient(), + mapper: mgr.GetRESTMapper(), + } + k8sClient, err := clientset.NewForConfig(mgr.GetConfig()) + if err != nil { + return err + } + Finder.discoveryClient = k8sClient.Discovery() + scaleKindResolver := scaleclient.NewDiscoveryScaleKindResolver(Finder.discoveryClient) + Finder.scaleNamespacer, err = scaleclient.NewForConfig(mgr.GetConfig(), Finder.mapper, dynamic.LegacyAPIPathResolverFunc, scaleKindResolver) + if err != nil { + return err + } + return nil +} + // ScaleAndSelector is used to return (controller, scale, selector) fields from the // controller finder functions. type ScaleAndSelector struct { @@ -42,8 +68,6 @@ type ScaleAndSelector struct { Selector *metav1.LabelSelector // metadata Metadata metav1.ObjectMeta - // template labels - TempLabels map[string]string } type ControllerReference struct { @@ -63,12 +87,10 @@ type PodControllerFinder func(ref ControllerReference, namespace string) (*Scale type ControllerFinder struct { client.Client -} -func NewControllerFinder(c client.Client) *ControllerFinder { - return &ControllerFinder{ - Client: c, - } + mapper meta.RESTMapper + scaleNamespacer scaleclient.ScalesGetter + discoveryClient discovery.DiscoveryInterface } func (r *ControllerFinder) GetExpectedScaleForPods(pods []*corev1.Pod) (int32, error) { @@ -123,19 +145,7 @@ func (r *ControllerFinder) GetScaleAndSelectorForRef(apiVersion, kind, ns, name func (r *ControllerFinder) Finders() []PodControllerFinder { return []PodControllerFinder{r.getPodReplicationController, r.getPodDeployment, r.getPodReplicaSet, - r.getPodStatefulSet, r.getPodKruiseCloneSet, r.getPodKruiseStatefulSet} -} - -func IsValidGroupVersionKind(apiVersion, kind string) bool { - for _, gvk := range validWorkloadList { - valid, err := verifyGroupKind(apiVersion, kind, gvk) - if err != nil { - return false - } else if valid { - return true - } - } - return false + r.getPodStatefulSet, r.getPodKruiseCloneSet, r.getPodKruiseStatefulSet, r.getScaleController} } var ( @@ -173,6 +183,7 @@ func (r *ControllerFinder) getPodReplicaSet(ref ControllerReference, namespace s } return r.getPodDeployment(refSs, namespace) } + return &ScaleAndSelector{ Scale: *(replicaSet.Spec.Replicas), Selector: replicaSet.Spec.Selector, @@ -182,8 +193,7 @@ func (r *ControllerFinder) getPodReplicaSet(ref ControllerReference, namespace s Name: replicaSet.Name, UID: replicaSet.UID, }, - Metadata: replicaSet.ObjectMeta, - TempLabels: replicaSet.Spec.Template.Labels, + Metadata: replicaSet.ObjectMeta, }, nil } @@ -238,8 +248,7 @@ func (r *ControllerFinder) getPodStatefulSet(ref ControllerReference, namespace Name: statefulSet.Name, UID: statefulSet.UID, }, - Metadata: statefulSet.ObjectMeta, - TempLabels: statefulSet.Spec.Template.Labels, + Metadata: statefulSet.ObjectMeta, }, nil } @@ -271,8 +280,7 @@ func (r *ControllerFinder) getPodDeployment(ref ControllerReference, namespace s Name: deployment.Name, UID: deployment.UID, }, - Metadata: deployment.ObjectMeta, - TempLabels: deployment.Spec.Template.Labels, + Metadata: deployment.ObjectMeta, }, nil } @@ -295,16 +303,14 @@ func (r *ControllerFinder) getPodReplicationController(ref ControllerReference, return nil, nil } return &ScaleAndSelector{ - Scale: *(rc.Spec.Replicas), - Selector: &metav1.LabelSelector{MatchLabels: rc.Spec.Selector}, + Scale: *(rc.Spec.Replicas), ControllerReference: ControllerReference{ APIVersion: rc.APIVersion, Kind: rc.Kind, Name: rc.Name, UID: rc.UID, }, - Metadata: rc.ObjectMeta, - TempLabels: rc.Spec.Template.Labels, + Metadata: rc.ObjectMeta, }, nil } @@ -337,8 +343,7 @@ func (r *ControllerFinder) getPodKruiseCloneSet(ref ControllerReference, namespa Name: cloneSet.Name, UID: cloneSet.UID, }, - Metadata: cloneSet.ObjectMeta, - TempLabels: cloneSet.Spec.Template.Labels, + Metadata: cloneSet.ObjectMeta, }, nil } @@ -372,8 +377,53 @@ func (r *ControllerFinder) getPodKruiseStatefulSet(ref ControllerReference, name Name: ss.Name, UID: ss.UID, }, - Metadata: ss.ObjectMeta, - TempLabels: ss.Spec.Template.Labels, + Metadata: ss.ObjectMeta, + }, nil +} + +func (r *ControllerFinder) getScaleController(ref ControllerReference, namespace string) (*ScaleAndSelector, error) { + if isValidGroupVersionKind(ref.APIVersion, ref.Kind) { + return nil, nil + } + gv, err := schema.ParseGroupVersion(ref.APIVersion) + if err != nil { + return nil, err + } + gk := schema.GroupKind{ + Group: gv.Group, + Kind: ref.Kind, + } + + mapping, err := r.mapper.RESTMapping(gk, gv.Version) + if err != nil { + return nil, err + } + gr := mapping.Resource.GroupResource() + scale, err := r.scaleNamespacer.Scales(namespace).Get(context.TODO(), gr, ref.Name, metav1.GetOptions{}) + if err != nil { + if errors.IsNotFound(err) { + // TODO, implementsScale + return nil, nil + } + return nil, err + } + if ref.UID != "" && scale.UID != ref.UID { + return nil, nil + } + selector, err := metav1.ParseToLabelSelector(scale.Status.Selector) + if err != nil { + return nil, err + } + return &ScaleAndSelector{ + Scale: scale.Spec.Replicas, + ControllerReference: ControllerReference{ + APIVersion: ref.APIVersion, + Kind: ref.Kind, + Name: ref.Name, + UID: scale.UID, + }, + Metadata: scale.ObjectMeta, + Selector: selector, }, nil } @@ -384,3 +434,15 @@ func verifyGroupKind(apiVersion, kind string, gvk schema.GroupVersionKind) (bool } return gv.Group == gvk.Group && kind == gvk.Kind, nil } + +func isValidGroupVersionKind(apiVersion, kind string) bool { + for _, gvk := range validWorkloadList { + valid, err := verifyGroupKind(apiVersion, kind, gvk) + if err != nil { + return false + } else if valid { + return true + } + } + return false +} diff --git a/pkg/util/controllerfinder/pods_finder.go b/pkg/util/controllerfinder/pods_finder.go index 258a13ec08..40e5d2e224 100644 --- a/pkg/util/controllerfinder/pods_finder.go +++ b/pkg/util/controllerfinder/pods_finder.go @@ -48,37 +48,38 @@ func (r *ControllerFinder) GetPodsForRef(apiVersion, kind, ns, name string, acti } workloadReplicas = *rs.Spec.Replicas workloadUIDs = append(workloadUIDs, rs.UID) - // Deployment, get the corresponding ReplicaSet UID - case ControllerKindDep.Kind: - rss, err := r.getReplicaSetsForDeployment(apiVersion, kind, ns, name) + // statefulset, rc, cloneSet + case ControllerKindSS.Kind, ControllerKindRC.Kind, ControllerKruiseKindCS.Kind, ControllerKruiseKindSS.Kind: + obj, err := r.GetScaleAndSelectorForRef(apiVersion, kind, ns, name, "") if err != nil { return nil, -1, err - } - if len(rss) == 0 { + } else if obj == nil { return nil, 0, nil } + workloadReplicas = obj.Scale + workloadUIDs = append(workloadUIDs, obj.UID) + // Deployment, Deployment-like workload, and other workload + default: obj, err := r.GetScaleAndSelectorForRef(apiVersion, kind, ns, name, "") if err != nil { return nil, -1, err - } - if obj == nil { + } else if obj == nil { return nil, 0, nil } workloadReplicas = obj.Scale - for _, rs := range rss { - workloadUIDs = append(workloadUIDs, rs.UID) - } - // others, e.g. rc, cloneset, statefulset... - default: - obj, err := r.GetScaleAndSelectorForRef(apiVersion, kind, ns, name, "") + // try to get replicaSets + rss, err := r.getReplicaSetsForDeployment(apiVersion, kind, ns, name) if err != nil { return nil, -1, err } - if obj == nil { - return nil, 0, nil + + if len(rss) == 0 { + workloadUIDs = append(workloadUIDs, obj.UID) + } else { + for _, rs := range rss { + workloadUIDs = append(workloadUIDs, rs.UID) + } } - workloadReplicas = obj.Scale - workloadUIDs = append(workloadUIDs, obj.UID) } // List all Pods owned by workload UID. @@ -106,12 +107,7 @@ func (r *ControllerFinder) GetPodsForRef(apiVersion, kind, ns, name string, acti } func (r *ControllerFinder) getReplicaSetsForDeployment(apiVersion, kind, ns, name string) ([]appsv1.ReplicaSet, error) { - targetRef := ControllerReference{ - APIVersion: apiVersion, - Kind: kind, - Name: name, - } - scaleNSelector, err := r.getPodDeployment(targetRef, ns) + scaleNSelector, err := r.GetScaleAndSelectorForRef(apiVersion, kind, ns, name, "") if err != nil || scaleNSelector == nil { return nil, err } diff --git a/pkg/webhook/pod/mutating/pod_create_update_handler.go b/pkg/webhook/pod/mutating/pod_create_update_handler.go index 3b1d73653a..1f332ba58a 100644 --- a/pkg/webhook/pod/mutating/pod_create_update_handler.go +++ b/pkg/webhook/pod/mutating/pod_create_update_handler.go @@ -111,7 +111,7 @@ var _ inject.Client = &PodCreateHandler{} // InjectClient injects the client into the PodCreateHandler func (h *PodCreateHandler) InjectClient(c client.Client) error { h.Client = c - h.finder = controllerfinder.NewControllerFinder(c) + h.finder = controllerfinder.Finder return nil } diff --git a/pkg/webhook/pod/mutating/pod_unavailable_budget.go b/pkg/webhook/pod/mutating/pod_unavailable_budget.go index 1d2e9d20b9..1dae0deb1d 100644 --- a/pkg/webhook/pod/mutating/pod_unavailable_budget.go +++ b/pkg/webhook/pod/mutating/pod_unavailable_budget.go @@ -20,9 +20,9 @@ import ( "context" "github.com/openkruise/kruise/pkg/control/pubcontrol" + "github.com/openkruise/kruise/pkg/controller/podunavailablebudget" admissionv1 "k8s.io/api/admission/v1" corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/klog/v2" "sigs.k8s.io/controller-runtime/pkg/webhook/admission" ) @@ -33,18 +33,7 @@ func (h *PodCreateHandler) pubMutatingPod(ctx context.Context, req admission.Req req.AdmissionRequest.Resource.Resource != "pods" { return nil } - ref := metav1.GetControllerOf(pod) - if ref == nil { - return nil - } - workload, err := h.finder.GetScaleAndSelectorForRef(ref.APIVersion, ref.Kind, pod.Namespace, ref.Name, "") - if err != nil { - return err - } else if workload == nil { - return nil - } - // fetch pub for workload - pub, err := pubcontrol.GetPubForWorkload(h.Client, workload) + pub, err := podunavailablebudget.GetPubForPod(h.Client, pod) if err != nil { return err } else if pub == nil { diff --git a/pkg/webhook/pod/validating/pod_create_update_handler.go b/pkg/webhook/pod/validating/pod_create_update_handler.go index 1fd914c676..cfb1fa7b16 100644 --- a/pkg/webhook/pod/validating/pod_create_update_handler.go +++ b/pkg/webhook/pod/validating/pod_create_update_handler.go @@ -90,7 +90,7 @@ var _ inject.Client = &PodCreateHandler{} // InjectClient injects the client into the PodCreateHandler func (h *PodCreateHandler) InjectClient(c client.Client) error { h.Client = c - h.finders = controllerfinder.NewControllerFinder(c) + h.finders = controllerfinder.Finder h.pubControl = pubcontrol.NewPubControl(c) return nil } diff --git a/pkg/webhook/pod/validating/pod_unavailable_budget.go b/pkg/webhook/pod/validating/pod_unavailable_budget.go index 3ef521f968..e88e36664b 100644 --- a/pkg/webhook/pod/validating/pod_unavailable_budget.go +++ b/pkg/webhook/pod/validating/pod_unavailable_budget.go @@ -18,12 +18,15 @@ package validating import ( "context" + "strings" + policyv1alpha1 "github.com/openkruise/kruise/apis/policy/v1alpha1" "github.com/openkruise/kruise/pkg/control/pubcontrol" admissionv1 "k8s.io/api/admission/v1" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/sets" "k8s.io/apiserver/pkg/util/dryrun" "k8s.io/klog/v2" "k8s.io/kubernetes/pkg/apis/policy" @@ -140,9 +143,10 @@ func (p *PodCreateHandler) podUnavailableBudgetValidatingPod(ctx context.Context pub, err := p.pubControl.GetPubForPod(newPod) if err != nil { return false, "", err - } - // if there is no matching PodUnavailableBudget, just return true - if pub == nil { + // if there is no matching PodUnavailableBudget, just return true + } else if pub == nil { + return true, "", nil + } else if !isNeedPubProtection(pub, pubcontrol.Operation(req.Operation)) { return true, "", nil } @@ -162,3 +166,22 @@ func (p *PodCreateHandler) podUnavailableBudgetValidatingPod(ctx context.Context return pubcontrol.PodUnavailableBudgetValidatePod(p.Client, p.pubControl, pub, newPod, pubcontrol.Operation(req.Operation), dryRun) } + +func isNeedPubProtection(pub *policyv1alpha1.PodUnavailableBudget, operation pubcontrol.Operation) bool { + operationValue, ok := pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] + if !ok { + return true + } + operations := sets.NewString() + for _, o := range strings.Split(operationValue, ",") { + operations.Insert(o) + } + // update operation + if operation == pubcontrol.UpdateOperation && operations.Has(pubcontrol.UpdateOperation) { + return true + // delete, eviction operation + } else if operation != pubcontrol.UpdateOperation && operations.Has(pubcontrol.DeleteOperation) { + return true + } + return false +} diff --git a/pkg/webhook/pod/validating/pod_unavailable_budget_test.go b/pkg/webhook/pod/validating/pod_unavailable_budget_test.go index 4290b21ca7..a203257293 100644 --- a/pkg/webhook/pod/validating/pod_unavailable_budget_test.go +++ b/pkg/webhook/pod/validating/pod_unavailable_budget_test.go @@ -56,8 +56,9 @@ var ( Kind: "PodUnavailableBudget", }, ObjectMeta: metav1.ObjectMeta{ - Namespace: "default", - Name: "pub-test", + Namespace: "default", + Name: "pub-test", + Annotations: map[string]string{}, }, Spec: policyv1alpha1.PodUnavailableBudgetSpec{ Selector: &metav1.LabelSelector{ @@ -421,6 +422,50 @@ func TestValidateUpdatePodForPub(t *testing.T) { return pubStatus }, }, + { + name: "valid update pod, pub feature-gate annotation, allow", + oldPod: func() *corev1.Pod { + pod := podDemo.DeepCopy() + return pod + }, + newPod: func() *corev1.Pod { + pod := podDemo.DeepCopy() + pod.Spec.Containers[0].Image = "nginx:1.18" + return pod + }, + pub: func() *policyv1alpha1.PodUnavailableBudget { + pub := pubDemo.DeepCopy() + pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] = "DELETE" + return pub + }, + expectAllow: true, + expectPubStatus: func() *policyv1alpha1.PodUnavailableBudgetStatus { + pubStatus := pubDemo.Status.DeepCopy() + return pubStatus + }, + }, + { + name: "valid update pod, pub feature-gate annotation, reject", + oldPod: func() *corev1.Pod { + pod := podDemo.DeepCopy() + return pod + }, + newPod: func() *corev1.Pod { + pod := podDemo.DeepCopy() + pod.Spec.Containers[0].Image = "nginx:1.18" + return pod + }, + pub: func() *policyv1alpha1.PodUnavailableBudget { + pub := pubDemo.DeepCopy() + pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] = "UPDATE" + return pub + }, + expectAllow: false, + expectPubStatus: func() *policyv1alpha1.PodUnavailableBudgetStatus { + pubStatus := pubDemo.Status.DeepCopy() + return pubStatus + }, + }, } for _, cs := range cases { @@ -663,6 +708,27 @@ func TestValidateDeletePodForPub(t *testing.T) { return pubStatus }, }, + { + name: "delete pod, pub feature-gate annotation, allow", + deletion: func() *metav1.DeleteOptions { + return &metav1.DeleteOptions{} + }, + newPod: func() *corev1.Pod { + podIn := podDemo.DeepCopy() + return podIn + }, + pub: func() *policyv1alpha1.PodUnavailableBudget { + pub := pubDemo.DeepCopy() + pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] = "UPDATE" + return pub + }, + subresource: "", + expectAllow: true, + expectPubStatus: func() *policyv1alpha1.PodUnavailableBudgetStatus { + pubStatus := pubDemo.Status.DeepCopy() + return pubStatus + }, + }, { name: "delete pod, allow", newPod: func() *corev1.Pod { diff --git a/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go b/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go index cc68f5f78c..c97cb509fa 100644 --- a/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go +++ b/pkg/webhook/podunavailablebudget/validating/pub_create_update_handler.go @@ -21,6 +21,7 @@ import ( "fmt" "net/http" "reflect" + "strings" policyv1alpha1 "github.com/openkruise/kruise/apis/policy/v1alpha1" "github.com/openkruise/kruise/pkg/features" @@ -82,8 +83,18 @@ func (h *PodUnavailableBudgetCreateUpdateHandler) Handle(ctx context.Context, re } func (h *PodUnavailableBudgetCreateUpdateHandler) validatingPodUnavailableBudgetFn(obj, old *policyv1alpha1.PodUnavailableBudget) field.ErrorList { + // validate pub.annotations + allErrs := field.ErrorList{} + if operationsValue, ok := obj.Annotations[policyv1alpha1.PubProtectOperationAnnotation]; ok { + operations := strings.Split(operationsValue, ",") + for _, operation := range operations { + if operation != string(admissionv1.Update) && operation != string(admissionv1.Delete) { + allErrs = append(allErrs, field.InternalError(field.NewPath("metadata"), fmt.Errorf("annotation[%s] is invalid", policyv1alpha1.PubProtectOperationAnnotation))) + } + } + } //validate Pub.Spec - allErrs := validatePodUnavailableBudgetSpec(obj, field.NewPath("spec")) + allErrs = append(allErrs, validatePodUnavailableBudgetSpec(obj, field.NewPath("spec"))...) // when operation is update, validating whether old and new pub conflict if old != nil { allErrs = append(allErrs, validateUpdatePubConflict(obj, old, field.NewPath("spec"))...) diff --git a/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go b/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go index 4fcd3872d3..3af4cf3afd 100644 --- a/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go +++ b/pkg/webhook/podunavailablebudget/validating/pub_validating_test.go @@ -43,8 +43,9 @@ var ( Kind: "PodUnavailableBudget", }, ObjectMeta: metav1.ObjectMeta{ - Namespace: "default", - Name: "pub-test", + Namespace: "default", + Name: "pub-test", + Annotations: map[string]string{}, }, Spec: policyv1alpha1.PodUnavailableBudgetSpec{ Selector: &metav1.LabelSelector{ @@ -135,6 +136,28 @@ func TestValidatingPub(t *testing.T) { }, expectErrList: 1, }, + { + name: "invalid pub feature-gate annotation", + pub: func() *policyv1alpha1.PodUnavailableBudget { + pub := pubDemo.DeepCopy() + pub.Spec.Selector = nil + pub.Spec.MinAvailable = nil + pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] = "xxxxx" + return pub + }, + expectErrList: 1, + }, + { + name: "valid pub feature-gate annotation", + pub: func() *policyv1alpha1.PodUnavailableBudget { + pub := pubDemo.DeepCopy() + pub.Spec.Selector = nil + pub.Spec.MinAvailable = nil + pub.Annotations[policyv1alpha1.PubProtectOperationAnnotation] = "DELETE" + return pub + }, + expectErrList: 0, + }, } decoder, _ := admission.NewDecoder(scheme) diff --git a/test/e2e/apps/daemonset.go b/test/e2e/apps/daemonset.go index f06b673a00..88c31eeb73 100644 --- a/test/e2e/apps/daemonset.go +++ b/test/e2e/apps/daemonset.go @@ -55,7 +55,6 @@ var _ = SIGDescribe("DaemonSet", func() { */ framework.ConformanceIt("should run and stop simple daemon", func() { label := map[string]string{framework.DaemonSetNameLabel: dsName} - ginkgo.By(fmt.Sprintf("Creating simple DaemonSet %q", dsName)) ds, err := tester.CreateDaemonSet(tester.NewDaemonSet(dsName, label, WebserverImage, appsv1alpha1.DaemonSetUpdateStrategy{})) gomega.Expect(err).NotTo(gomega.HaveOccurred())