From 5e6d7bdf04ecf4ef264004df55497b0803a094e9 Mon Sep 17 00:00:00 2001
From: "liheng.zms" <liheng.zms@alibaba-inc.com>
Date: Tue, 4 Jul 2023 16:10:08 +0800
Subject: [PATCH] auto patch webhook objectSelector label on workload

---
 config/webhook/patch_manifests.yaml      | 22 ++++++++++++++++++
 pkg/controller/rollout/rollout_status.go | 29 ++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/config/webhook/patch_manifests.yaml b/config/webhook/patch_manifests.yaml
index 84d7959b..5022bcab 100644
--- a/config/webhook/patch_manifests.yaml
+++ b/config/webhook/patch_manifests.yaml
@@ -8,6 +8,26 @@ webhooks:
       matchExpressions:
         - key: rollouts.kruise.io/workload-type
           operator: Exists
+  - name: mcloneset.kb.io
+    objectSelector:
+      matchExpressions:
+        - key: rollouts.kruise.io/workload-type
+          operator: Exists
+  - name: mdaemonset.kb.io
+    objectSelector:
+      matchExpressions:
+        - key: rollouts.kruise.io/workload-type
+          operator: Exists
+  - name: mstatefulset.kb.io
+    objectSelector:
+      matchExpressions:
+        - key: rollouts.kruise.io/workload-type
+          operator: Exists
+  - name: madvancedstatefulset.kb.io
+    objectSelector:
+      matchExpressions:
+        - key: rollouts.kruise.io/workload-type
+          operator: Exists
   - name: mdeployment.kb.io
     objectSelector:
       matchExpressions:
@@ -15,3 +35,5 @@ webhooks:
           operator: NotIn
           values:
             - controller-manager
+        - key: rollouts.kruise.io/workload-type
+          operator: Exists
diff --git a/pkg/controller/rollout/rollout_status.go b/pkg/controller/rollout/rollout_status.go
index f1bd8c10..7ca66059 100644
--- a/pkg/controller/rollout/rollout_status.go
+++ b/pkg/controller/rollout/rollout_status.go
@@ -25,6 +25,7 @@ import (
 	"github.com/openkruise/rollouts/api/v1alpha1"
 	"github.com/openkruise/rollouts/pkg/util"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/client-go/util/retry"
@@ -66,8 +67,32 @@ func (r *RolloutReconciler) calculateRolloutStatus(rollout *v1alpha1.Rollout) (r
 		klog.Infof("rollout(%s/%s) workload not found, and reset status be Initial", rollout.Namespace, rollout.Name)
 		return false, newStatus, nil
 	}
-	klog.V(5).Infof("rollout(%s/%s) workload(%s)", rollout.Namespace, rollout.Name, util.DumpJSON(workload))
-	// todo, patch workload webhook labels
+	klog.V(5).Infof("rollout(%s/%s) fetch workload(%s)", rollout.Namespace, rollout.Name, util.DumpJSON(workload))
+	// Patch objectSelector in workload labels[rollouts.kruise.io/workload-type],
+	// then rollout only webhook the workload which contains labels[rollouts.kruise.io/workload-type].
+	var workloadType util.WorkloadType
+	switch workload.Kind {
+	case util.ControllerKruiseKindCS.Kind:
+		workloadType = util.CloneSetType
+	case util.ControllerKindDep.Kind:
+		workloadType = util.DeploymentType
+	case util.ControllerKindSts.Kind:
+		workloadType = util.StatefulSetType
+	case util.ControllerKruiseKindDS.Kind:
+		workloadType = util.DaemonSetType
+	}
+	if workload.Annotations[util.WorkloadTypeLabel] == "" && workloadType != "" {
+		workloadGVK := schema.FromAPIVersionAndKind(workload.APIVersion, workload.Kind)
+		obj := util.GetEmptyWorkloadObject(workloadGVK)
+		obj.SetNamespace(workload.Namespace)
+		obj.SetName(workload.Name)
+		body := fmt.Sprintf(`{"metadata":{"labels":{"%s":"%s"}}}`, util.WorkloadTypeLabel, workloadType)
+		if err = r.Patch(context.TODO(), obj, client.RawPatch(types.MergePatchType, []byte(body))); err != nil {
+			klog.Errorf("rollout(%s/%s) patch workload(%s) failed: %s", rollout.Namespace, rollout.Name, workload.Name, err.Error())
+			return false, nil, err
+		}
+		klog.Infof("rollout(%s/%s) patch workload(%s) labels[%s] success", rollout.Namespace, rollout.Name, workload.Name, util.WorkloadTypeLabel)
+	}
 	// workload status generation is not equal to workload.generation
 	if !workload.IsStatusConsistent {
 		klog.Infof("rollout(%s/%s) workload status is inconsistent, then wait a moment", rollout.Namespace, rollout.Name)