Add configurable taint effect (#539)

Add TAINT_EFFECT env var and --taint-effect flag to configure the node taint effect
Defaults to NoSchedule, garbage values will log a warning and default to NoSchedule.

closes  #273

Author: hamishforbes

pkg/config/config.go

@@ -67,6 +67,8 @@ const (
 	metadataTriesDefault                    = 3
 	cordonOnly                              = "CORDON_ONLY"
 	taintNode                               = "TAINT_NODE"
+	taintEffectDefault                      = "NoSchedule"
+	taintEffect                             = "TAINT_EFFECT"
 	jsonLoggingConfigKey                    = "JSON_LOGGING"
 	jsonLoggingDefault                      = false
 	logLevelConfigKey                       = "LOG_LEVEL"
@@ -123,6 +125,7 @@ type Config struct {
 	MetadataTries                    int
 	CordonOnly                       bool
 	TaintNode                        bool
+	TaintEffect                      string
 	JsonLogging                      bool
 	LogLevel                         string
 	UptimeFromFile                   string
@@ -175,6 +178,7 @@ func ParseCliArgs() (config Config, err error) {
 	flag.IntVar(&config.MetadataTries, "metadata-tries", getIntEnv(metadataTriesConfigKey, metadataTriesDefault), "The number of times to try requesting metadata. If you would like 2 retries, set metadata-tries to 3.")
 	flag.BoolVar(&config.CordonOnly, "cordon-only", getBoolEnv(cordonOnly, false), "If true, nodes will be cordoned but not drained when an interruption event occurs.")
 	flag.BoolVar(&config.TaintNode, "taint-node", getBoolEnv(taintNode, false), "If true, nodes will be tainted when an interruption event occurs.")
+	flag.StringVar(&config.TaintEffect, "taint-effect", getEnv(taintEffect, taintEffectDefault), "Sets the effect when a node is tainted.")
 	flag.BoolVar(&config.JsonLogging, "json-logging", getBoolEnv(jsonLoggingConfigKey, jsonLoggingDefault), "If true, use JSON-formatted logs instead of human readable logs.")
 	flag.StringVar(&config.LogLevel, "log-level", getEnv(logLevelConfigKey, logLevelDefault), "Sets the log level (INFO, DEBUG, or ERROR)")
 	flag.StringVar(&config.UptimeFromFile, "uptime-from-file", getEnv(uptimeFromFileConfigKey, uptimeFromFileDefault), "If specified, read system uptime from the file path (useful for testing).")
@@ -249,6 +253,7 @@ func (c Config) PrintJsonConfigArgs() {
 		Int("metadata_tries", c.MetadataTries).
 		Bool("cordon_only", c.CordonOnly).
 		Bool("taint_node", c.TaintNode).
+		Str("taint_effect", c.TaintEffect).
 		Bool("json_logging", c.JsonLogging).
 		Str("log_level", c.LogLevel).
 		Str("webhook_proxy", c.WebhookProxy).
@@ -292,6 +297,7 @@ func (c Config) PrintHumanConfigArgs() {
 			"\tmetadata-tries: %d,\n"+
 			"\tcordon-only: %t,\n"+
 			"\ttaint-node: %t,\n"+
+			"\ttaint-effect: %s,\n"+
 			"\tjson-logging: %t,\n"+
 			"\tlog-level: %s,\n"+
 			"\twebhook-proxy: %s,\n"+
@@ -326,6 +332,7 @@ func (c Config) PrintHumanConfigArgs() {
 		c.MetadataTries,
 		c.CordonOnly,
 		c.TaintNode,
+		c.TaintEffect,
 		c.JsonLogging,
 		c.LogLevel,
 		c.WebhookProxy,

pkg/node/node.go

@@ -310,7 +310,7 @@ func (n Node) TaintSpotItn(nodeName string, eventID string) error {
 		eventID = eventID[:maxTaintValueLength]
 	}
 
-	return addTaint(k8sNode, n, SpotInterruptionTaint, eventID, corev1.TaintEffectNoSchedule)
+	return addTaint(k8sNode, n, SpotInterruptionTaint, eventID)
 }
 
 // TaintASGLifecycleTermination adds the spot termination notice taint onto a node
@@ -328,7 +328,7 @@ func (n Node) TaintASGLifecycleTermination(nodeName string, eventID string) erro
 		eventID = eventID[:maxTaintValueLength]
 	}
 
-	return addTaint(k8sNode, n, ASGLifecycleTerminationTaint, eventID, corev1.TaintEffectNoSchedule)
+	return addTaint(k8sNode, n, ASGLifecycleTerminationTaint, eventID)
 }
 
 // TaintRebalanceRecommendation adds the rebalance recommendation notice taint onto a node
@@ -346,7 +346,7 @@ func (n Node) TaintRebalanceRecommendation(nodeName string, eventID string) erro
 		eventID = eventID[:maxTaintValueLength]
 	}
 
-	return addTaint(k8sNode, n, RebalanceRecommendationTaint, eventID, corev1.TaintEffectNoSchedule)
+	return addTaint(k8sNode, n, RebalanceRecommendationTaint, eventID)
 }
 
 // LogPods logs all the pod names on a node
@@ -388,7 +388,7 @@ func (n Node) TaintScheduledMaintenance(nodeName string, eventID string) error {
 		eventID = eventID[:maxTaintValueLength]
 	}
 
-	return addTaint(k8sNode, n, ScheduledMaintenanceTaint, eventID, corev1.TaintEffectNoSchedule)
+	return addTaint(k8sNode, n, ScheduledMaintenanceTaint, eventID)
 }
 
 // RemoveNTHTaints removes NTH-specific taints from a node
@@ -540,7 +540,22 @@ func jsonPatchEscape(value string) string {
 	return strings.Replace(value, "/", "~1", -1)
 }
 
-func addTaint(node *corev1.Node, nth Node, taintKey string, taintValue string, effect corev1.TaintEffect) error {
+func getTaintEffect(effect string) corev1.TaintEffect {
+	switch effect {
+	case "PreferNoSchedule":
+		return corev1.TaintEffectPreferNoSchedule
+	case "NoExecute":
+		return corev1.TaintEffectNoExecute
+	default:
+		log.Warn().Msgf("Unknown taint effect: %s", effect)
+		fallthrough
+	case "NoSchedule":
+		return corev1.TaintEffectNoSchedule
+	}
+}
+
+func addTaint(node *corev1.Node, nth Node, taintKey string, taintValue string) error {
+	effect := getTaintEffect(nth.nthConfig.TaintEffect)
 	if nth.nthConfig.DryRun {
 		log.Info().Msgf("Would have added taint (%s=%s:%s) to node %s, but dry-run flag was set", taintKey, taintValue, effect, nth.nthConfig.NodeName)
 		return nil