From de7c9f27bd59f05af59c6572b097d316532704ec Mon Sep 17 00:00:00 2001
From: shaojiang <shaojiang.deng@apecloud.com>
Date: Tue, 20 Jun 2023 18:51:02 +0800
Subject: [PATCH] fix: add tolerations for cronjob (#3881)

(cherry picked from commit 903226ed2946b68321b78c67090e1871c0a3f651)
---
 .../dataprotection/backuppolicy_controller.go       |  5 +++++
 .../dataprotection/backuppolicy_controller_test.go  | 13 +++++++++++++
 controllers/dataprotection/cue/cronjob.cue          |  8 ++++++++
 controllers/dataprotection/type.go                  |  2 ++
 4 files changed, 28 insertions(+)

diff --git a/controllers/dataprotection/backuppolicy_controller.go b/controllers/dataprotection/backuppolicy_controller.go
index 1e1f41b0eb3..e3572120532 100644
--- a/controllers/dataprotection/backuppolicy_controller.go
+++ b/controllers/dataprotection/backuppolicy_controller.go
@@ -282,6 +282,10 @@ func (r *BackupPolicyReconciler) buildCronJob(
 	if err != nil {
 		return nil, err
 	}
+	tolerationPodSpec := corev1.PodSpec{}
+	if err = addTolerations(&tolerationPodSpec); err != nil {
+		return nil, err
+	}
 	var ttl metav1.Duration
 	if backupPolicy.Spec.Retention != nil && backupPolicy.Spec.Retention.TTL != nil {
 		ttl = metav1.Duration{Duration: dataprotectionv1alpha1.ToDuration(backupPolicy.Spec.Retention.TTL)}
@@ -298,6 +302,7 @@ func (r *BackupPolicyReconciler) buildCronJob(
 		ServiceAccount:   viper.GetString("KUBEBLOCKS_SERVICEACCOUNT_NAME"),
 		MgrNamespace:     viper.GetString(constant.CfgKeyCtrlrMgrNS),
 		Image:            viper.GetString(constant.KBToolsImage),
+		Tolerations:      &tolerationPodSpec,
 	}
 	backupPolicyOptionsByte, err := json.Marshal(options)
 	if err != nil {
diff --git a/controllers/dataprotection/backuppolicy_controller_test.go b/controllers/dataprotection/backuppolicy_controller_test.go
index f4dafc97d5f..f4ffe7a18e1 100644
--- a/controllers/dataprotection/backuppolicy_controller_test.go
+++ b/controllers/dataprotection/backuppolicy_controller_test.go
@@ -117,6 +117,12 @@ var _ = Describe("Backup Policy Controller", func() {
 
 		BeforeEach(func() {
 			viper.Set(constant.CfgKeyCtrlrMgrNS, mgrNamespace)
+			viper.Set(constant.CfgKeyCtrlrMgrAffinity,
+				"{\"nodeAffinity\":{\"preferredDuringSchedulingIgnoredDuringExecution\":[{\"preference\":{\"matchExpressions\":[{\"key\":\"kb-controller\",\"operator\":\"In\",\"values\":[\"true\"]}]},\"weight\":100}]}}")
+			viper.Set(constant.CfgKeyCtrlrMgrTolerations,
+				"[{\"key\":\"key1\", \"operator\": \"Exists\", \"effect\": \"NoSchedule\"}]")
+			viper.Set(constant.CfgKeyCtrlrMgrNodeSelector, "{\"beta.kubernetes.io/arch\":\"amd64\"}")
+
 			By("By creating a backupTool")
 			backupTool := testapps.CreateCustomizedObj(&testCtx, "backup/backuptool.yaml",
 				&dpv1alpha1.BackupTool{}, testapps.RandomizedObjName())
@@ -125,6 +131,9 @@ var _ = Describe("Backup Policy Controller", func() {
 
 		AfterEach(func() {
 			viper.SetDefault(constant.CfgKeyCtrlrMgrNS, testCtx.DefaultNamespace)
+			viper.Set(constant.CfgKeyCtrlrMgrAffinity, "")
+			viper.Set(constant.CfgKeyCtrlrMgrTolerations, "")
+			viper.Set(constant.CfgKeyCtrlrMgrNodeSelector, "")
 		})
 
 		Context("creates a backup policy", func() {
@@ -152,6 +161,10 @@ var _ = Describe("Backup Policy Controller", func() {
 				})).Should(Succeed())
 				Eventually(testapps.CheckObj(&testCtx, getCronjobKey(dpv1alpha1.BackupTypeDataFile), func(g Gomega, fetched *batchv1.CronJob) {
 					g.Expect(fetched.Spec.Schedule).To(Equal(defaultSchedule))
+					g.Expect(fetched.Spec.JobTemplate.Spec.Template.Spec.Tolerations).ShouldNot(BeEmpty())
+					g.Expect(fetched.Spec.JobTemplate.Spec.Template.Spec.NodeSelector).ShouldNot(BeEmpty())
+					g.Expect(fetched.Spec.JobTemplate.Spec.Template.Spec.Affinity).ShouldNot(BeNil())
+					g.Expect(fetched.Spec.JobTemplate.Spec.Template.Spec.Affinity.NodeAffinity).ShouldNot(BeNil())
 				})).Should(Succeed())
 			})
 			It("limit backups to 1", func() {
diff --git a/controllers/dataprotection/cue/cronjob.cue b/controllers/dataprotection/cue/cronjob.cue
index 8a2fe9956d0..5680ce7b801 100644
--- a/controllers/dataprotection/cue/cronjob.cue
+++ b/controllers/dataprotection/cue/cronjob.cue
@@ -26,6 +26,11 @@ options: {
 	ttl:              string
 	serviceAccount:   string
 	image:            string
+	tolerations: {
+		tolerations: [...]
+		affinity: {...}
+		nodeSelector: {...}
+	}
 }
 
 cronjob: {
@@ -47,6 +52,9 @@ cronjob: {
 		jobTemplate: spec: template: spec: {
 			restartPolicy:      "Never"
 			serviceAccountName: options.serviceAccount
+			affinity: options.tolerations.affinity
+			tolerations: options.tolerations.tolerations
+			nodeSelector: options.tolerations.nodeSelector
 			containers: [{
 				name:            "backup-policy"
 				image:           options.image
diff --git a/controllers/dataprotection/type.go b/controllers/dataprotection/type.go
index 26d4ed014b4..851330c9a7e 100644
--- a/controllers/dataprotection/type.go
+++ b/controllers/dataprotection/type.go
@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/spf13/viper"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -72,4 +73,5 @@ type backupPolicyOptions struct {
 	TTL              metav1.Duration `json:"ttl,omitempty"`
 	ServiceAccount   string          `json:"serviceAccount"`
 	Image            string          `json:"image"`
+	Tolerations      *corev1.PodSpec `json:"tolerations"`
 }