pod mutate patch affinity

Signed-off-by: min.jia <jiamin13579@126.com>

pkg/webhooks/admission/pods/mutate/mutate_pod.go

@@ -122,6 +122,11 @@ func createPatch(pod *v1.Pod) ([]byte, error) {
 			patch = append(patch, *patchLabel)
 		}
 
+		patchAffinity := patchAffinity(pod, resourceGroup)
+		if patchAffinity != nil {
+			patch = append(patch, *patchAffinity)
+		}
+
 		patchToleration := patchTaintToleration(pod, resourceGroup)
 		if patchToleration != nil {
 			patch = append(patch, *patchToleration)
@@ -156,6 +161,28 @@ func patchLabels(pod *v1.Pod, resGroupConfig wkconfig.ResGroupConfig) *patchOper
 	return &patchOperation{Op: "add", Path: "/spec/nodeSelector", Value: nodeSelector}
 }
 
+// patchAffinity patch affinity
+func patchAffinity(pod *v1.Pod, resGroupConfig wkconfig.ResGroupConfig) *patchOperation {
+	if resGroupConfig.Affinity == "" {
+		return nil
+	}
+
+	if pod.Spec.Affinity != nil {
+		klog.V(5).Infof("pod affinity exist: %s", pod.Name)
+		return nil
+	}
+
+	var affinity v1.Affinity
+	err := json.Unmarshal([]byte(resGroupConfig.Affinity), &affinity)
+	if err != nil {
+		fmt.Println("Failed to unmarshal JSON:", err)
+		klog.V(3).Infof("Failed to unmarshal JSON: %s", err)
+		return nil
+	}
+
+	return &patchOperation{Op: "add", Path: "/spec/affinity", Value: affinity}
+}
+
 // patchTaintToleration patch taint toleration
 func patchTaintToleration(pod *v1.Pod, resGroupConfig wkconfig.ResGroupConfig) *patchOperation {
 	if len(resGroupConfig.Tolerations) == 0 {

pkg/webhooks/admission/pods/mutate/mutate_pod_test.go

@@ -28,6 +28,9 @@ import (
 )
 
 func TestMutatePods(t *testing.T) {
+	affinityJsonStr := `{"nodeAffinity":{"requiredDuringSchedulingIgnoredDuringExecution":{"nodeSelectorTerms":[{"matchExpressions":[{"key":"kubernetes.io/os","operator":"In","values":["linux"]}]}]}}}`
+	var affinity v1.Affinity
+	json.Unmarshal([]byte(affinityJsonStr), &affinity)
 
 	admissionConfigData := &webconfig.AdmissionConfiguration{
 		ResGroupsConfig: []webconfig.ResGroupConfig{
@@ -48,6 +51,7 @@ func TestMutatePods(t *testing.T) {
 						Effect:   v1.TaintEffectNoSchedule,
 					},
 				},
+				Affinity: affinityJsonStr,
 				Labels: map[string]string{
 					"volcano.sh/nodetype": "management",
 				},
@@ -105,6 +109,11 @@ func TestMutatePods(t *testing.T) {
 						"volcano.sh/nodetype": "management",
 					},
 				},
+				{
+					Op:    "add",
+					Path:  "/spec/affinity",
+					Value: affinity,
+				},
 				{
 					Op:   "add",
 					Path: "/spec/tolerations",

pkg/webhooks/config/config.go

@@ -42,6 +42,7 @@ type ResGroupConfig struct {
 	SchedulerName string            `yaml:"schedulerName"`
 	Tolerations   []v1.Toleration   `yaml:"tolerations"`
 	Labels        map[string]string `yaml:"labels"`
+	Affinity      string            `yaml:"affinity"`
 }
 
 // AdmissionConfiguration defines the configuration of admission.