add patch for UnitedDeployment (openkruise#1266)

generate manifests and fix goimports fix goimports patch to PodTemplateSpec add ut for uniteddeployment patch Signed-off-by: chengleqi <[email protected]>
MarkLux · Jun 5, 2023 · 6fda363 · 6fda363
1 parent 19240cf
commit 6fda363
Show file tree

Hide file tree

Showing 15 changed files with 981 additions and 11 deletions.
diff --git a/apis/apps/v1alpha1/uniteddeployment_types.go b/apis/apps/v1alpha1/uniteddeployment_types.go
@@ -21,6 +21,7 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/intstr"
 )
 
@@ -187,6 +188,13 @@ type Subset struct {
 	// Controller will try to keep all the subsets with nil replicas have average pods.
 	// +optional
 	Replicas *intstr.IntOrString `json:"replicas,omitempty"`
+	// Patch indicates patching to the templateSpec.
+	// Patch takes precedence over other fields
+	// If the Patch also modifies the Replicas, NodeSelectorTerm or Tolerations, use value in the Patch
+	// +optional
+	// +kubebuilder:pruning:PreserveUnknownFields
+	// +kubebuilder:validation:Schemaless
+	Patch runtime.RawExtension `json:"patch,omitempty"`
 }
 
 // UnitedDeploymentStatus defines the observed state of UnitedDeployment.

diff --git a/apis/apps/v1alpha1/well_known_labels.go b/apis/apps/v1alpha1/well_known_labels.go
@@ -15,6 +15,8 @@ const (
 
 	// ImagePreDownloadIgnoredKey indicates the images of this revision have been ignored to pre-download
 	ImagePreDownloadIgnoredKey = "apps.kruise.io/image-predownload-ignored"
+	// AnnotationSubsetPatchKey indicates the patch for every subset
+	AnnotationSubsetPatchKey = "apps.kruise.io/subset-patch"
 )
 
 // Sidecar container environment variable definitions which are used to enable SidecarTerminator to take effect on the sidecar container.

diff --git a/apis/apps/v1alpha1/zz_generated.deepcopy.go b/apis/apps/v1alpha1/zz_generated.deepcopy.go
diff --git a/config/crd/bases/apps.kruise.io_uniteddeployments.yaml b/config/crd/bases/apps.kruise.io_uniteddeployments.yaml
@@ -1056,6 +1056,12 @@ spec:
                                 type: object
                               type: array
                           type: object
+                        patch:
+                          description: Patch indicates patching to the templateSpec.
+                            Patch takes precedence over other fields If the Patch
+                            also modifies the Replicas, NodeSelectorTerm or Tolerations,
+                            use value in the Patch
+                          x-kubernetes-preserve-unknown-fields: true
                         replicas:
                           anyOf:
                           - type: integer

diff --git a/pkg/controller/uniteddeployment/adapter/advanced_statefulset_adapter.go b/pkg/controller/uniteddeployment/adapter/advanced_statefulset_adapter.go
@@ -18,8 +18,12 @@ package adapter
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
+	"k8s.io/apimachinery/pkg/util/strategicpatch"
+	"k8s.io/klog/v2"
+
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -173,6 +177,26 @@ func (a *AdvancedStatefulSetAdapter) ApplySubsetTemplate(ud *alpha1.UnitedDeploy
 
 	attachNodeAffinity(&set.Spec.Template.Spec, subSetConfig)
 	attachTolerations(&set.Spec.Template.Spec, subSetConfig)
+	if subSetConfig.Patch.Raw != nil {
+		TemplateSpecBytes, _ := json.Marshal(set.Spec.Template)
+		modified, err := strategicpatch.StrategicMergePatch(TemplateSpecBytes, subSetConfig.Patch.Raw, &corev1.PodTemplateSpec{})
+		if err != nil {
+			klog.Errorf("failed to merge patch raw %s", subSetConfig.Patch.Raw)
+			return err
+		}
+		patchedTemplateSpec := corev1.PodTemplateSpec{}
+		if err = json.Unmarshal(modified, &patchedTemplateSpec); err != nil {
+			klog.Errorf("failed to unmarshal %s to podTemplateSpec", modified)
+			return err
+		}
+
+		set.Spec.Template = patchedTemplateSpec
+		klog.V(2).Infof("AdvancedStatefulSet [%s/%s] was patched successfully: %s", set.Namespace, set.GenerateName, subSetConfig.Patch.Raw)
+	}
+	if set.Annotations == nil {
+		set.Annotations = make(map[string]string)
+	}
+	set.Annotations[alpha1.AnnotationSubsetPatchKey] = string(subSetConfig.Patch.Raw)
 
 	return nil
 }

diff --git a/pkg/controller/uniteddeployment/adapter/cloneset_adapter.go b/pkg/controller/uniteddeployment/adapter/cloneset_adapter.go
@@ -2,8 +2,12 @@ package adapter
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
+	"k8s.io/apimachinery/pkg/util/strategicpatch"
+	"k8s.io/klog/v2"
+
 	alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
 	"github.com/openkruise/kruise/pkg/util"
 	"github.com/openkruise/kruise/pkg/util/refmanager"
@@ -134,7 +138,26 @@ func (a *CloneSetAdapter) ApplySubsetTemplate(ud *alpha1.UnitedDeployment, subse
 
 	attachNodeAffinity(&set.Spec.Template.Spec, subSetConfig)
 	attachTolerations(&set.Spec.Template.Spec, subSetConfig)
+	if subSetConfig.Patch.Raw != nil {
+		TemplateSpecBytes, _ := json.Marshal(set.Spec.Template)
+		modified, err := strategicpatch.StrategicMergePatch(TemplateSpecBytes, subSetConfig.Patch.Raw, &corev1.PodTemplateSpec{})
+		if err != nil {
+			klog.Errorf("failed to merge patch raw %s", subSetConfig.Patch.Raw)
+			return err
+		}
+		patchedTemplateSpec := corev1.PodTemplateSpec{}
+		if err = json.Unmarshal(modified, &patchedTemplateSpec); err != nil {
+			klog.Errorf("failed to unmarshal %s to podTemplateSpec", modified)
+			return err
+		}
 
+		set.Spec.Template = patchedTemplateSpec
+		klog.V(2).Infof("CloneSet [%s/%s] was patched successfully: %s", set.Namespace, set.GenerateName, subSetConfig.Patch.Raw)
+	}
+	if set.Annotations == nil {
+		set.Annotations = make(map[string]string)
+	}
+	set.Annotations[alpha1.AnnotationSubsetPatchKey] = string(subSetConfig.Patch.Raw)
 	return nil
 }
 

diff --git a/pkg/controller/uniteddeployment/adapter/deployment_adapter.go b/pkg/controller/uniteddeployment/adapter/deployment_adapter.go
@@ -18,8 +18,12 @@ package adapter
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
+	"k8s.io/apimachinery/pkg/util/strategicpatch"
+	"k8s.io/klog/v2"
+
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -146,6 +150,27 @@ func (a *DeploymentAdapter) ApplySubsetTemplate(ud *alpha1.UnitedDeployment, sub
 	attachNodeAffinity(&set.Spec.Template.Spec, subSetConfig)
 	attachTolerations(&set.Spec.Template.Spec, subSetConfig)
 
+	if subSetConfig.Patch.Raw != nil {
+		TemplateSpecBytes, _ := json.Marshal(set.Spec.Template)
+		modified, err := strategicpatch.StrategicMergePatch(TemplateSpecBytes, subSetConfig.Patch.Raw, &corev1.PodTemplateSpec{})
+		if err != nil {
+			klog.Errorf("failed to merge patch raw %s", subSetConfig.Patch.Raw)
+			return err
+		}
+		patchedTemplateSpec := corev1.PodTemplateSpec{}
+		if err = json.Unmarshal(modified, &patchedTemplateSpec); err != nil {
+			klog.Errorf("failed to unmarshal %s to podTemplateSpec", modified)
+			return err
+		}
+
+		set.Spec.Template = patchedTemplateSpec
+		klog.V(2).Infof("Deployment [%s/%s] was patched successfully: %s", set.Namespace, set.GenerateName, subSetConfig.Patch.Raw)
+	}
+	if set.Annotations == nil {
+		set.Annotations = make(map[string]string)
+	}
+	set.Annotations[alpha1.AnnotationSubsetPatchKey] = string(subSetConfig.Patch.Raw)
+
 	return nil
 }
 

diff --git a/pkg/controller/uniteddeployment/adapter/statefulset_adapter.go b/pkg/controller/uniteddeployment/adapter/statefulset_adapter.go
@@ -18,8 +18,11 @@ package adapter
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
+	"k8s.io/apimachinery/pkg/util/strategicpatch"
+
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -157,6 +160,26 @@ func (a *StatefulSetAdapter) ApplySubsetTemplate(ud *alpha1.UnitedDeployment, su
 
 	attachNodeAffinity(&set.Spec.Template.Spec, subSetConfig)
 	attachTolerations(&set.Spec.Template.Spec, subSetConfig)
+	if subSetConfig.Patch.Raw != nil {
+		TemplateSpecBytes, _ := json.Marshal(set.Spec.Template)
+		modified, err := strategicpatch.StrategicMergePatch(TemplateSpecBytes, subSetConfig.Patch.Raw, &corev1.PodTemplateSpec{})
+		if err != nil {
+			klog.Errorf("failed to merge patch raw %s", subSetConfig.Patch.Raw)
+			return err
+		}
+		patchedTemplateSpec := corev1.PodTemplateSpec{}
+		if err = json.Unmarshal(modified, &patchedTemplateSpec); err != nil {
+			klog.Errorf("failed to unmarshal %s to podTemplateSpec", modified)
+			return err
+		}
+
+		set.Spec.Template = patchedTemplateSpec
+		klog.V(2).Infof("StatefulSet [%s/%s] was patched successfully: %s", set.Namespace, set.GenerateName, subSetConfig.Patch.Raw)
+	}
+	if set.Annotations == nil {
+		set.Annotations = make(map[string]string)
+	}
+	set.Annotations[alpha1.AnnotationSubsetPatchKey] = string(subSetConfig.Patch.Raw)
 
 	return nil
 }

diff --git a/pkg/controller/uniteddeployment/subset.go b/pkg/controller/uniteddeployment/subset.go
@@ -52,6 +52,13 @@ type SubsetUpdateStrategy struct {
 	Partition int32
 }
 
+// SubsetUpdate stores the subset field that may need to be updated
+type SubsetUpdate struct {
+	Replicas  int32
+	Partition int32
+	Patch     string
+}
+
 // ResourceRef stores the Subset resource it represents.
 type ResourceRef struct {
 	Resources []metav1.Object

diff --git a/pkg/controller/uniteddeployment/uniteddeployment_controller.go b/pkg/controller/uniteddeployment/uniteddeployment_controller.go
@@ -222,9 +222,10 @@ func (r *ReconcileUnitedDeployment) Reconcile(_ context.Context, request reconci
 	}
 
 	nextPartitions := calcNextPartitions(instance, nextReplicas)
-	klog.V(4).Infof("Get UnitedDeployment %s/%s next partition %v", instance.Namespace, instance.Name, nextPartitions)
+	nextUpdate := getNextUpdate(instance, nextReplicas, nextPartitions)
+	klog.V(4).Infof("Get UnitedDeployment %s/%s next update %v", instance.Namespace, instance.Name, nextUpdate)
 
-	newStatus, err := r.manageSubsets(instance, nameToSubset, nextReplicas, nextPartitions, currentRevision, updatedRevision, subsetType)
+	newStatus, err := r.manageSubsets(instance, nameToSubset, nextUpdate, currentRevision, updatedRevision, subsetType)
 	if err != nil {
 		klog.Errorf("Fail to update UnitedDeployment %s/%s: %s", instance.Namespace, instance.Name, err)
 		r.recorder.Event(instance.DeepCopy(), corev1.EventTypeWarning, fmt.Sprintf("Failed%s", eventTypeSubsetsUpdate), err.Error())
@@ -273,6 +274,19 @@ func calcNextPartitions(ud *appsv1alpha1.UnitedDeployment, nextReplicas *map[str
 	return &partitions
 }
 
+func getNextUpdate(ud *appsv1alpha1.UnitedDeployment, nextReplicas *map[string]int32, nextPartitions *map[string]int32) map[string]SubsetUpdate {
+	next := make(map[string]SubsetUpdate)
+	for _, subset := range ud.Spec.Topology.Subsets {
+		t := SubsetUpdate{}
+		t.Replicas = (*nextReplicas)[subset.Name]
+		t.Partition = (*nextPartitions)[subset.Name]
+		t.Patch = string(subset.Patch.Raw)
+
+		next[subset.Name] = t
+	}
+	return next
+}
+
 func (r *ReconcileUnitedDeployment) deleteDupSubset(ud *appsv1alpha1.UnitedDeployment, nameToSubsets map[string][]*Subset, control ControlInterface) (*map[string]*Subset, error) {
 	nameToSubset := map[string]*Subset{}
 	for name, subsets := range nameToSubsets {