Add admission plugin for PodGroup to add finalizer to every new object

Signed-off-by: helayoty <heelayot@microsoft.com>

pkg/controlplane/apiserver/samples/generic/server/admission_test.go

@@ -29,6 +29,7 @@ import (
 	"k8s.io/kubernetes/plugin/pkg/admission/podtopologylabels"
 	podpriority "k8s.io/kubernetes/plugin/pkg/admission/priority"
 	"k8s.io/kubernetes/plugin/pkg/admission/runtimeclass"
+	"k8s.io/kubernetes/plugin/pkg/admission/scheduling/podgroupprotection"
 	"k8s.io/kubernetes/plugin/pkg/admission/security/podsecurity"
 	"k8s.io/kubernetes/plugin/pkg/admission/storage/persistentvolume/resize"
 	"k8s.io/kubernetes/plugin/pkg/admission/storage/storageclass/setdefault"
@@ -40,6 +41,7 @@ var intentionallyOffPlugins = sets.New[string](
 	setdefault.PluginName,                   // DefaultStorageClass
 	resize.PluginName,                       // PersistentVolumeClaimResize
 	storageobjectinuseprotection.PluginName, // StorageObjectInUseProtection
+	podgroupprotection.PluginName,           // PodGroupProtection
 	podpriority.PluginName,                  // Priority
 	nodetaint.PluginName,                    // TaintNodesByCondition
 	runtimeclass.PluginName,                 // RuntimeClass

pkg/kubeapiserver/options/plugins.go

@@ -53,6 +53,7 @@ import (
 	"k8s.io/kubernetes/plugin/pkg/admission/podtopologylabels"
 	podpriority "k8s.io/kubernetes/plugin/pkg/admission/priority"
 	"k8s.io/kubernetes/plugin/pkg/admission/runtimeclass"
+	"k8s.io/kubernetes/plugin/pkg/admission/scheduling/podgroupprotection"
 	"k8s.io/kubernetes/plugin/pkg/admission/security/podsecurity"
 	"k8s.io/kubernetes/plugin/pkg/admission/serviceaccount"
 	"k8s.io/kubernetes/plugin/pkg/admission/storage/persistentvolume/resize"
@@ -89,6 +90,7 @@ var AllOrderedPlugins = []string{
 	extendedresourcetoleration.PluginName,   // ExtendedResourceToleration
 	setdefault.PluginName,                   // DefaultStorageClass
 	storageobjectinuseprotection.PluginName, // StorageObjectInUseProtection
+	podgroupprotection.PluginName,           // PodGroupProtection
 	gc.PluginName,                           // OwnerReferencesPermissionEnforcement
 	resize.PluginName,                       // PersistentVolumeClaimResize
 	runtimeclass.PluginName,                 // RuntimeClass
@@ -148,6 +150,7 @@ func RegisterAllAdmissionPlugins(plugins *admission.Plugins) {
 	setdefault.Register(plugins)
 	resize.Register(plugins)
 	storageobjectinuseprotection.Register(plugins)
+	podgroupprotection.Register(plugins)
 	certapproval.Register(plugins)
 	certsigning.Register(plugins)
 	ctbattest.Register(plugins)
@@ -170,6 +173,7 @@ func DefaultOffAdmissionPlugins() sets.Set[string] {
 		validatingwebhook.PluginName,            // ValidatingAdmissionWebhook
 		resourcequota.PluginName,                // ResourceQuota
 		storageobjectinuseprotection.PluginName, // StorageObjectInUseProtection
+		podgroupprotection.PluginName,           // PodGroupProtection
 		podpriority.PluginName,                  // Priority
 		nodetaint.PluginName,                    // TaintNodesByCondition
 		runtimeclass.PluginName,                 // RuntimeClass

plugin/pkg/admission/scheduling/podgroupprotection/admission.go

@@ -0,0 +1,101 @@
+/*
+Copyright The Kubernetes 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 podgroupprotection
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"slices"
+
+	"k8s.io/apiserver/pkg/admission"
+	apiserveradmission "k8s.io/apiserver/pkg/admission/initializer"
+	"k8s.io/component-base/featuregate"
+	"k8s.io/klog/v2"
+	schedulingapi "k8s.io/kubernetes/pkg/apis/scheduling"
+	"k8s.io/kubernetes/pkg/features"
+)
+
+const (
+	PluginName = "PodGroupProtection"
+)
+
+// Register registers the plugin.
+func Register(plugins *admission.Plugins) {
+	plugins.Register(PluginName, func(config io.Reader) (admission.Interface, error) {
+		return newPlugin(), nil
+	})
+}
+
+type podGroupProtectionPlugin struct {
+	*admission.Handler
+	enabled               bool
+	inspectedFeatureGates bool
+}
+
+var _ admission.MutationInterface = &podGroupProtectionPlugin{}
+var _ apiserveradmission.WantsFeatures = &podGroupProtectionPlugin{}
+
+func newPlugin() *podGroupProtectionPlugin {
+	return &podGroupProtectionPlugin{
+		Handler: admission.NewHandler(admission.Create),
+	}
+}
+
+func (p *podGroupProtectionPlugin) InspectFeatureGates(featureGates featuregate.FeatureGate) {
+	p.enabled = featureGates.Enabled(features.GenericWorkload)
+	p.inspectedFeatureGates = true
+}
+
+func (p *podGroupProtectionPlugin) ValidateInitialization() error {
+	if !p.inspectedFeatureGates {
+		return fmt.Errorf("feature gates not inspected")
+	}
+	return nil
+}
+
+var podGroupResource = schedulingapi.Resource("podgroups")
+
+// Admit stamps the PodGroupProtectionFinalizer on every newly created PodGroup
+// so that it cannot be deleted while pods still reference it.
+// The finalizer is removed by the PodGroupProtection controller when the
+// PodGroup is no longer in use.
+func (p *podGroupProtectionPlugin) Admit(ctx context.Context, a admission.Attributes, o admission.ObjectInterfaces) error {
+	if !p.enabled {
+		return nil
+	}
+
+	if a.GetResource().GroupResource() != podGroupResource {
+		return nil
+	}
+	if len(a.GetSubresource()) != 0 {
+		return nil
+	}
+
+	pg, ok := a.GetObject().(*schedulingapi.PodGroup)
+	if !ok {
+		return nil
+	}
+
+	if slices.Contains(pg.Finalizers, schedulingapi.PodGroupProtectionFinalizer) {
+		return nil
+	}
+
+	klog.V(4).InfoS("Adding protection finalizer to PodGroup", "podGroup", klog.KObj(pg))
+	pg.Finalizers = append(pg.Finalizers, schedulingapi.PodGroupProtectionFinalizer)
+	return nil
+}

plugin/pkg/admission/scheduling/podgroupprotection/admission_test.go

@@ -0,0 +1,111 @@
+/*
+Copyright The Kubernetes 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 podgroupprotection
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apiserver/pkg/admission"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
+	schedulingapi "k8s.io/kubernetes/pkg/apis/scheduling"
+	"k8s.io/kubernetes/pkg/features"
+	"k8s.io/utils/dump"
+)
+
+func TestAdmit(t *testing.T) {
+	pg := &schedulingapi.PodGroup{
+		TypeMeta: metav1.TypeMeta{Kind: "PodGroup"},
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "my-podgroup",
+			Namespace: "default",
+		},
+	}
+
+	pgWithFinalizer := pg.DeepCopy()
+	pgWithFinalizer.Finalizers = []string{schedulingapi.PodGroupProtectionFinalizer}
+
+	tests := []struct {
+		name           string
+		enabled        bool
+		resource       schema.GroupVersionResource
+		object         runtime.Object
+		expectedObject runtime.Object
+		namespace      string
+	}{
+		{
+			name:           "podgroup create with plugin enabled, add finalizer",
+			enabled:        true,
+			resource:       schedulingapi.SchemeGroupVersion.WithResource("podgroups"),
+			object:         pg,
+			expectedObject: pgWithFinalizer,
+			namespace:      pg.Namespace,
+		},
+		{
+			name:           "podgroup finalizer already exists, no new finalizer",
+			enabled:        true,
+			resource:       schedulingapi.SchemeGroupVersion.WithResource("podgroups"),
+			object:         pgWithFinalizer,
+			expectedObject: pgWithFinalizer,
+			namespace:      pgWithFinalizer.Namespace,
+		},
+		{
+			name:           "podgroup create with plugin disabled, no finalizer added",
+			enabled:        false,
+			resource:       schedulingapi.SchemeGroupVersion.WithResource("podgroups"),
+			object:         pg,
+			expectedObject: pg,
+			namespace:      pg.Namespace,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.GenericWorkload, test.enabled)
+
+			ctrl := newPlugin()
+			ctrl.InspectFeatureGates(utilfeature.DefaultFeatureGate)
+
+			obj := test.object.DeepCopyObject()
+			attrs := admission.NewAttributesRecord(
+				obj,
+				obj.DeepCopyObject(),
+				schema.GroupVersionKind{},
+				test.namespace,
+				"foo",
+				test.resource,
+				"",
+				admission.Create,
+				&metav1.CreateOptions{},
+				false,
+				nil,
+			)
+
+			if err := ctrl.Admit(context.TODO(), attrs, nil); err != nil {
+				t.Errorf("got unexpected error: %v", err)
+			}
+			if !reflect.DeepEqual(test.expectedObject, obj) {
+				t.Errorf("Expected object:\n%s\ngot:\n%s", dump.Pretty(test.expectedObject), dump.Pretty(obj))
+			}
+		})
+	}
+}