fix(agent): respect taint-toleration

Ref: Longhorn-5614

Signed-off-by: Chin-Ya Huang <chin-ya.huang@suse.com>

cmd/manager.go

@@ -51,6 +51,7 @@ func init() {
 	managerCmd.PersistentFlags().StringVar(&sbm.ImageName, "image-name", os.Getenv("SUPPORT_BUNDLE_IMAGE"), "The support bundle image")
 	managerCmd.PersistentFlags().StringVar(&sbm.ImagePullPolicy, "image-pull-policy", os.Getenv("SUPPORT_BUNDLE_IMAGE_PULL_POLICY"), "Pull policy of the support bundle image")
 	managerCmd.PersistentFlags().StringVar(&sbm.NodeSelector, "node-selector", os.Getenv("SUPPORT_BUNDLE_NODE_SELECTOR"), "NodeSelector of agent DaemonSet. e.g., key1=value1,key2=value2")
+	managerCmd.PersistentFlags().StringVar(&sbm.TaintToleration, "taint-toleration", os.Getenv("SUPPORT_BUNDLE_TAINT_TOLERATION"), "Toleration of agent DaemonSet. e.g., key1=value1:NoSchedule,key2=value2:NoSchedule")
 	managerCmd.PersistentFlags().StringVar(&sbm.RegistrySecret, "registry-secret", os.Getenv("SUPPORT_BUNDLE_REGISTRY_SECRET"), "The registry secret for image pull")
 	managerCmd.PersistentFlags().StringVar(&sbm.SpecifyCollector, "specify-collector", os.Getenv("SUPPORT_BUNDLE_COLLECTOR"), "Execute specify collector script. e.g., longhorn")
 	managerCmd.PersistentFlags().StringSliceVar(&sbm.ExcludeResourceList, "exclude-resources", getEnvStringSlice("SUPPORT_BUNDLE_EXCLUDE_RESOURCES"), "List of resources to exclude. e.g., settings.harvesterhci.io,secrets")

pkg/manager/agent.go

@@ -68,11 +68,7 @@ func (a *AgentDaemonSet) Create(image string, managerURL string) error {
 				},
 				Spec: corev1.PodSpec{
 					NodeSelector: a.sbm.getNodeSelector(),
-					Tolerations: []corev1.Toleration{
-						{
-							Operator: corev1.TolerationOpExists,
-						},
-					},
+					Tolerations:  a.sbm.getTaintToleration(),
 					Containers: []corev1.Container{
 						{
 							Name:            "agent",

pkg/manager/manager.go

@@ -36,6 +36,7 @@ type SupportBundleManager struct {
 	KubeConfig      string
 	PodNamespace    string
 	NodeSelector    string
+	TaintToleration string
 	RegistrySecret  string
 
 	ExcludeResources    []schema.GroupResource
@@ -378,3 +379,73 @@ func (m *SupportBundleManager) getNodeSelector() map[string]string {
 	}
 	return nodeSelector
 }
+
+func (m *SupportBundleManager) getTaintToleration() []v1.Toleration {
+	taintToleration := []v1.Toleration{}
+
+	m.TaintToleration = strings.ReplaceAll(m.TaintToleration, " ", "")
+	if m.TaintToleration == "" {
+		return taintToleration
+	}
+
+	tolerationList := strings.Split(m.TaintToleration, ",")
+	for _, toleration := range tolerationList {
+		toleration, err := parseToleration(toleration)
+		if err != nil {
+			logrus.WithError(err).Warnf("Invalid toleration: %s", toleration)
+			continue
+		}
+		taintToleration = append(taintToleration, *toleration)
+	}
+
+	if len(taintToleration) == 0 {
+		return []v1.Toleration{{Operator: v1.TolerationOpExists}}
+	}
+	return taintToleration
+}
+
+func parseToleration(taintToleration string) (*v1.Toleration, error) {
+	toleration := &v1.Toleration{}
+
+	// The schema should be `key=value:effect` or `key:effect`
+	parts := strings.Split(taintToleration, ":")
+	if len(parts) != 2 {
+		return nil, fmt.Errorf("missing key/value and effect pair")
+	}
+
+	keyValue := parts[0]
+	key, value, operator, err := parseTolerationKeyAndValue(keyValue)
+	if err != nil {
+		return nil, err
+	}
+	toleration.Key = key
+	toleration.Value = value
+	toleration.Operator = operator
+	toleration.Effect, err = parseTaintEffect(parts[1])
+	if err != nil {
+		return nil, err
+	}
+	return toleration, nil
+}
+
+func parseTolerationKeyAndValue(keyValue string) (string, string, v1.TolerationOperator, error) {
+	// parse `key=value` or `key`
+	if strings.Contains(keyValue, "=") {
+		pair := strings.Split(keyValue, "=")
+		if len(pair) != 2 {
+			return "", "", v1.TolerationOpEqual, fmt.Errorf("invalid key/value pair: %v", keyValue)
+		}
+		return pair[0], pair[1], v1.TolerationOpEqual, nil
+	}
+	return keyValue, "", v1.TolerationOpExists, nil
+}
+
+func parseTaintEffect(effect string) (v1.TaintEffect, error) {
+	taintEffect := v1.TaintEffect(effect)
+	switch taintEffect {
+	case v1.TaintEffectNoExecute, v1.TaintEffectNoSchedule, v1.TaintEffectPreferNoSchedule, v1.TaintEffect(""):
+		return taintEffect, nil
+	default:
+		return "", fmt.Errorf("invalid effect: %v", effect)
+	}
+}