add resourrcecontext webhook and improve logic (#312)

* add resourrcecontext webhook

* remove other collasets podContext if nesscessary

* remove other collasets podContext if nesscessary

* reset manager

* reset manager

* add e2e

pkg/controllers/collaset/podcontext/podcontext.go

@@ -24,9 +24,9 @@ import (
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	appsv1alpha1 "kusionstack.io/kube-api/apps/v1alpha1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
-	appsv1alpha1 "kusionstack.io/kube-api/apps/v1alpha1"
 	"kusionstack.io/kuperator/pkg/controllers/collaset/utils"
 	"kusionstack.io/kuperator/pkg/controllers/utils/expectations"
 )
@@ -61,9 +61,11 @@ func AllocateID(c client.Client, instance *appsv1alpha1.CollaSet, defaultRevisio
 		detail := &podContext.Spec.Contexts[i]
 		if detail.Contains(OwnerContextKey, instance.Name) {
 			ownedIDs[detail.ID] = detail
+			existingIDs[detail.ID] = detail
+		} else if instance.Spec.ScaleStrategy.Context != "" {
+			// add other collaset podContexts only if context pool enabled
+			existingIDs[detail.ID] = detail
 		}
-
-		existingIDs[detail.ID] = detail
 	}
 
 	// if owner has enough ID, return
@@ -149,13 +151,16 @@ func doUpdatePodContext(c client.Client, instance client.Object, ownedIDs map[in
 	// store all IDs crossing all workload
 	existingIDs := map[int]*appsv1alpha1.ContextDetail{}
 
-	for i := range podContext.Spec.Contexts {
-		detail := podContext.Spec.Contexts[i]
-		if detail.Contains(OwnerContextKey, instance.GetName()) {
-			continue
+	// add other collaset podContexts only if context pool enabled
+	cls := instance.(*appsv1alpha1.CollaSet)
+	if cls.Spec.ScaleStrategy.Context != "" {
+		for i := range podContext.Spec.Contexts {
+			detail := podContext.Spec.Contexts[i]
+			if detail.Contains(OwnerContextKey, instance.GetName()) {
+				continue
+			}
+			existingIDs[detail.ID] = &detail
 		}
-
-		existingIDs[detail.ID] = &detail
 	}
 
 	for _, contextDetail := range ownedIDs {

pkg/webhook/server/generic/generic_webhooks.go

@@ -23,6 +23,7 @@ import (
 	"kusionstack.io/kuperator/pkg/webhook/server/generic/operationjob"
 	"kusionstack.io/kuperator/pkg/webhook/server/generic/persistentvolumeclaim"
 	"kusionstack.io/kuperator/pkg/webhook/server/generic/poddecoration"
+	"kusionstack.io/kuperator/pkg/webhook/server/generic/resourcecontext"
 
 	webhookdmission "kusionstack.io/kuperator/pkg/webhook/admission"
 	"kusionstack.io/kuperator/pkg/webhook/server/generic/pod"
@@ -60,4 +61,7 @@ func init() {
 
 	MutatingTypeHandlerMap["OperationJob"] = operationjob.NewMutatingHandler()
 	ValidatingTypeHandlerMap["OperationJob"] = operationjob.NewValidatingHandler()
+
+	MutatingTypeHandlerMap["ResourceContext"] = resourcecontext.NewMutatingHandler()
+	ValidatingTypeHandlerMap["ResourceContext"] = resourcecontext.NewValidatingHandler()
 }

pkg/webhook/server/generic/resourcecontext/resourcecontext_mutating_handler.go

@@ -0,0 +1,43 @@
+/*
+Copyright 2024 The KusionStack 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 resourcecontext
+
+import (
+	"context"
+
+	"sigs.k8s.io/controller-runtime/pkg/runtime/inject"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
+	"kusionstack.io/kuperator/pkg/utils/mixin"
+)
+
+var _ inject.Client = &MutatingHandler{}
+var _ admission.DecoderInjector = &MutatingHandler{}
+
+type MutatingHandler struct {
+	*mixin.WebhookHandlerMixin
+}
+
+func NewMutatingHandler() *MutatingHandler {
+	return &MutatingHandler{
+		WebhookHandlerMixin: mixin.NewWebhookHandlerMixin(),
+	}
+}
+
+func (h *MutatingHandler) Handle(ctx context.Context, req admission.Request) (resp admission.Response) {
+	return admission.Allowed("")
+}

pkg/webhook/server/generic/resourcecontext/resourcecontext_validating_handler.go

@@ -0,0 +1,43 @@
+/*
+Copyright 2024 The KusionStack 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 resourcecontext
+
+import (
+	"context"
+
+	"sigs.k8s.io/controller-runtime/pkg/runtime/inject"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
+	"kusionstack.io/kuperator/pkg/utils/mixin"
+)
+
+var _ inject.Client = &ValidatingHandler{}
+var _ admission.DecoderInjector = &ValidatingHandler{}
+
+type ValidatingHandler struct {
+	*mixin.WebhookHandlerMixin
+}
+
+func NewValidatingHandler() *ValidatingHandler {
+	return &ValidatingHandler{
+		WebhookHandlerMixin: mixin.NewWebhookHandlerMixin(),
+	}
+}
+
+func (h *ValidatingHandler) Handle(ctx context.Context, req admission.Request) (resp admission.Response) {
+	return admission.Allowed("")
+}

test/e2e/apps/collaset.go

@@ -38,6 +38,7 @@ import (
 	appsv1alpha1 "kusionstack.io/kube-api/apps/v1alpha1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
+	"kusionstack.io/kuperator/pkg/controllers/collaset/podcontext"
 	collasetutils "kusionstack.io/kuperator/pkg/controllers/collaset/utils"
 	"kusionstack.io/kuperator/pkg/controllers/utils/podopslifecycle"
 	"kusionstack.io/kuperator/pkg/utils"
@@ -1380,6 +1381,62 @@ var _ = SIGDescribe("CollaSet", func() {
 				pod.Finalizers = []string{}
 			})).NotTo(HaveOccurred())
 		})
+
+		framework.ConformanceIt("update with hacking resourceContexts", func() {
+			cls := tester.NewCollaSet("collaset-"+randStr, 3, appsv1alpha1.UpdateStrategy{})
+			cls.Spec.UpdateStrategy = appsv1alpha1.UpdateStrategy{
+				RollingUpdate: &appsv1alpha1.RollingUpdateCollaSetStrategy{
+					ByLabel: &appsv1alpha1.ByLabel{},
+				},
+			}
+			Expect(tester.CreateCollaSet(cls)).NotTo(HaveOccurred())
+
+			By("Wait for status replicas satisfied")
+			Eventually(func() error { return tester.ExpectedStatusReplicas(cls, 3, 3, 3, 3, 3) }, 30*time.Second, 3*time.Second).ShouldNot(HaveOccurred())
+
+			By("Hacking resourceContexts with different collaset")
+			currResourceContexts, err := tester.ListResourceContextsForCollaSet(cls)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(len(currResourceContexts[0].Spec.Contexts)).Should(BeEquivalentTo(3))
+			for i := range currResourceContexts[0].Spec.Contexts {
+				currResourceContexts[0].Spec.Contexts[i].Data[podcontext.OwnerContextKey] = "mock-collaset"
+			}
+			Expect(client.Update(context.Background(), currResourceContexts[0])).Should(BeNil())
+
+			By("Update image to nginxNew but pods are not updated")
+			Expect(tester.UpdateCollaSet(cls, func(cls *appsv1alpha1.CollaSet) {
+				cls.Spec.Template.Spec.Containers[0].Image = imageutils.GetE2EImage(imageutils.NginxNew)
+			})).NotTo(HaveOccurred())
+			Eventually(func() error { return tester.ExpectedStatusReplicas(cls, 3, 3, 3, 0, 3) }, 30*time.Second, 3*time.Second).ShouldNot(HaveOccurred())
+
+			By("Wait for CollaSet reconciled")
+			Eventually(func() bool {
+				if err := tester.GetCollaSet(cls); err != nil {
+					return false
+				}
+				return cls.Generation == cls.Status.ObservedGeneration
+			}, 10*time.Second, 3*time.Second).Should(Equal(true))
+
+			By("Label pod to trigger update")
+			pods, err := tester.ListPodsForCollaSet(cls)
+			Expect(err).NotTo(HaveOccurred())
+			podToUpdate := pods[0]
+			Expect(tester.UpdatePod(podToUpdate, func(pod *v1.Pod) {
+				pod.Labels[appsv1alpha1.CollaSetUpdateIndicateLabelKey] = "true"
+			})).NotTo(HaveOccurred())
+
+			By("Wait for update finished")
+			Eventually(func() error { return tester.ExpectedStatusReplicas(cls, 3, 3, 3, 1, 3) }, 30*time.Second, 3*time.Second).ShouldNot(HaveOccurred())
+
+			By("Check resourceContexts")
+			currResourceContexts, err = tester.ListResourceContextsForCollaSet(cls)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(len(currResourceContexts[0].Spec.Contexts)).Should(BeEquivalentTo(3))
+			for i := range currResourceContexts[0].Spec.Contexts {
+				Expect(currResourceContexts[0].Spec.Contexts[i].Data[podcontext.OwnerContextKey]).Should(BeEquivalentTo(cls.Name))
+			}
+		})
+
 	})
 
 	framework.KusionstackDescribe("CollaSet Replacing", func() {