Allow tolerations to be passed to helm jobs

main.go

@@ -53,6 +53,12 @@ func main() {
 				EnvVars:     []string{"DEFAULT_JOB_IMAGE"},
 				Destination: &config.DefaultJobImage,
 			},
+			&cli.StringFlag{
+				Name:        "job-tolerations",
+				Usage:       "JSON array of tolerations to apply to all jobs managing helm charts",
+				EnvVars:     []string{"JOB_TOLERATIONS"},
+				Destination: &config.JobTolerations,
+			},
 			&cli.StringFlag{
 				Name:        "job-cluster-role",
 				Value:       "cluster-admin",

pkg/cmd/cmd.go

@@ -1,9 +1,11 @@
 package cmd
 
 import (
+	"encoding/json"
 	"fmt"
 	"log"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -14,6 +16,7 @@ import (
 	"github.com/rancher/wrangler/v3/pkg/kubeconfig"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
@@ -36,6 +39,7 @@ type HelmController struct {
 	NodeName        string
 	JobClusterRole  string
 	DefaultJobImage string
+	JobTolerations  string
 	PprofPort       int
 }
 
@@ -78,11 +82,17 @@ func (hc *HelmController) Run(app *cli.Context) error {
 		return err
 	}
 
+	tolerations, err := parseTolerations(hc.JobTolerations)
+	if err != nil {
+		return fmt.Errorf("invalid --job-tolerations JSON: %w", err)
+	}
+
 	opts := common.Options{
 		Threadiness:     hc.Threads,
 		NodeName:        hc.NodeName,
 		JobClusterRole:  hc.JobClusterRole,
 		DefaultJobImage: hc.DefaultJobImage,
+		JobTolerations:  tolerations,
 	}
 
 	if err := opts.Validate(); err != nil {
@@ -110,3 +120,17 @@ func (hc *HelmController) GetNonInteractiveClientConfig() clientcmd.ClientConfig
 			ClusterInfo:     clientcmdapi.Cluster{Server: hc.MasterURL},
 		}, nil)
 }
+
+// parseTolerations takes the CLI string and parses it into a slice of corev1.Toleration objects
+func parseTolerations(raw string) ([]corev1.Toleration, error) {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return nil, nil
+	}
+
+	var tolerations []corev1.Toleration
+	if err := json.Unmarshal([]byte(raw), &tolerations); err != nil {
+		return nil, err
+	}
+	return tolerations, nil
+}

pkg/controllers/chart/chart.go

@@ -71,8 +71,8 @@ const (
 
 var (
 	commaRE              = regexp.MustCompile(`\\*,`)
-	deletePolicy         = metav1.DeletePropagationForeground
 	DefaultJobImage      = "rancher/klipper-helm:latest"
+	JobTolerations       []corev1.Toleration
 	DefaultFailurePolicy = FailurePolicyReinstall
 	defaultBackOffLimit  = ptr.To(int32(1000))
 	jobSettleTime        = time.Second * 3
@@ -863,6 +863,7 @@ func job(chart *v1.HelmChart, apiServerPort string) (*batch.Job, *corev1.Secret,
 	setDockerRegistrySecret(job, chart)
 	setRepoCAConfigMap(job, chart)
 	setSecurityContext(job, chart)
+	setTolerations(job)
 	valuesSecret := setValuesSecret(job, chart)
 	contentConfigMap := setContentConfigMap(job, chart)
 
@@ -1303,6 +1304,12 @@ func setSecurityContext(job *batch.Job, chart *v1.HelmChart) {
 	}
 }
 
+func setTolerations(job *batch.Job) {
+	if len(JobTolerations) > 0 {
+		job.Spec.Template.Spec.Tolerations = append(job.Spec.Template.Spec.Tolerations, JobTolerations...)
+	}
+}
+
 // chartSource returns a string describing the source of the chart:
 // chartContent, chart URL, or repo+version
 func chartSource(chart *v1.HelmChart) string {

pkg/controllers/chart/chart_test.go

@@ -181,6 +181,38 @@ func TestInstallJobImage(t *testing.T) {
 	assert.Equal("custom-job-image", job.Spec.Template.Spec.Containers[0].Image)
 }
 
+func TestInstallJobTolerations(t *testing.T) {
+	assert := assert.New(t)
+	chart := NewChart()
+	oldDefaultJobTolerations := JobTolerations
+	defer func() { JobTolerations = oldDefaultJobTolerations }()
+	JobTolerations = []corev1.Toleration{{
+		Key:      "custom-taint",
+		Operator: corev1.TolerationOpExists,
+		Effect:   corev1.TaintEffectNoSchedule,
+	}}
+
+	job, _, _ := job(chart, "6443")
+	assert.Contains(job.Spec.Template.Spec.Tolerations, JobTolerations[0])
+}
+
+func TestInstallJobBootstrapAndCustomTolerations(t *testing.T) {
+	assert := assert.New(t)
+	chart := NewChart()
+	chart.Spec.Bootstrap = true
+	oldDefaultJobTolerations := JobTolerations
+	defer func() { JobTolerations = oldDefaultJobTolerations }()
+	JobTolerations = []corev1.Toleration{{
+		Key:      "custom-taint",
+		Operator: corev1.TolerationOpExists,
+		Effect:   corev1.TaintEffectNoExecute,
+	}}
+
+	job, _, _ := job(chart, "6443")
+	assert.GreaterOrEqual(len(job.Spec.Template.Spec.Tolerations), len(JobTolerations)+1)
+	assert.Contains(job.Spec.Template.Spec.Tolerations, JobTolerations[0])
+}
+
 func TestInstallArgs(t *testing.T) {
 	assert := assert.New(t)
 	stringArgs := strings.Join(args(NewChart()), " ")

pkg/controllers/common/options.go

@@ -1,13 +1,18 @@
 package common
 
-import "fmt"
+import (
+	"fmt"
+
+	corev1 "k8s.io/api/core/v1"
+)
 
 // Options defines options that can be set on initializing the Helm Controller
 type Options struct {
 	Threadiness     int
 	NodeName        string
 	JobClusterRole  string
 	DefaultJobImage string
+	JobTolerations  []corev1.Toleration
 }
 
 func (opts Options) Validate() error {

pkg/controllers/controllers.go

@@ -82,6 +82,7 @@ func Register(ctx context.Context, systemNamespace, controllerName string, cfg c
 	if opts.DefaultJobImage != "" {
 		chart.DefaultJobImage = opts.DefaultJobImage
 	}
+	chart.JobTolerations = opts.JobTolerations
 
 	chart.Register(ctx,
 		systemNamespace,
@@ -108,6 +109,7 @@ func Register(ctx context.Context, systemNamespace, controllerName string, cfg c
 	logger.Info("Starting helm controller", "threads", opts.Threadiness)
 	logger.Info("Using cluster role for jobs managing helm charts", "jobClusterRole", opts.JobClusterRole)
 	logger.Info("Using default image for jobs managing helm charts", "defaultJobImage", chart.DefaultJobImage)
+	logger.Info("Using tolerations for jobs managing helm charts", "jobTolerationsCount", len(chart.JobTolerations))
 
 	if len(systemNamespace) == 0 {
 		systemNamespace = metav1.NamespaceSystem