Add --tolerations option to kots cli (#5075)

* Add --tolerations option to kots cli * use selected tolerations in kotsadm, minio, rqlite and distribution pods * improve test code * fix op names, begin unit test * fleshed out unit tests * quote invalid toleration part values * intermediary greps should not use '-q' * add 'all tolerations are present' log line --------- Co-authored-by: Jim Falgout <[email protected]>
replicatedhq · Jan 2, 2025 · 2b704db · 2b704db
1 parent d92aaab
commit 2b704db
Show file tree

Hide file tree

Showing 8 changed files with 212 additions and 2 deletions.
diff --git a/.github/workflows/build-test.yaml b/.github/workflows/build-test.yaml
@@ -30,12 +30,12 @@ jobs:
       - name: ok
         run: echo "yes"
 
-  # Use this to disable tests when iteratig on a specific test to save time
+  # Use this to disable tests when iterating on a specific test to save time
   enable-tests:
     runs-on: ubuntu-20.04
     steps:
       - name: ok
-        # change 0 to a positive interger to prevent all tests from running
+        # change 0 to a positive integer to prevent all tests from running
         run: exit 0
 
 
@@ -695,6 +695,8 @@ jobs:
             --additional-labels test.label/two=test.value.two \
             --additional-annotations testannotation1=testannotationvalue1 \
             --additional-annotations test.annotation/two=testannotation.value.two \
+            --tolerations test.com/role:Equal:core:NoSchedule \
+            --tolerations test.com/productid:Exists::NoSchedule \
             --kotsadm-tag 24h | tee output.txt
 
           if ! grep -q "The Kubernetes RBAC policy that the Admin Console is running with does not have access to complete the Preflight Checks. It's recommended that you run these manually before proceeding." output.txt; then
@@ -803,6 +805,43 @@ jobs:
           fi
           echo "additional pod labels and annotations are present"
 
+          echo "check that the kotsadm, minio and rqlite pods have the correct tolerations"
+          echo "Equal toleration"
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Equal' | grep 'NoSchedule' | grep -q 'test.com/role'; then
+              echo "kotsadm pod does not have Equal toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm-minio -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Equal' | grep 'NoSchedule' | grep -q 'test.com/role'; then
+              echo "kotsadm-minio pod does not have Equal toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm-minio -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm-rqlite -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Equal' | grep 'NoSchedule' | grep -q 'test.com/role'; then
+              echo "kotsadm-rqlite pod does not have Equal toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm-rqlite -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          echo "Equal tolerations are present"
+          echo "Exists toleration"
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Exists' | grep 'NoSchedule' | grep -q 'test.com/productid'; then
+              echo "kotsadm pod does not have Exists toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm-minio -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Exists' | grep 'NoSchedule' | grep -q 'test.com/productid'; then
+              echo "kotsadm-minio pod does not have Exists toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm-minio -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          if ! kubectl get pods -n "$APP_SLUG" -l app=kotsadm-rqlite -o jsonpath='{.items[0].spec.tolerations[*]}' | grep 'Exists' | grep 'NoSchedule' | grep -q 'test.com/productid'; then
+              echo "kotsadm-rqlite pod does not have Exists toleration"
+              kubectl get pods -n "$APP_SLUG" -l app=kotsadm-rqlite -o jsonpath='{.items[0].spec.tolerations[*]}'
+              exit 1
+          fi
+          echo "Exists tolerations are present"
+          echo "all tolerations are present"
+
       - name: Generate support bundle on failure
         if: failure()
         uses: ./.github/actions/generate-support-bundle

diff --git a/cmd/kots/cli/install.go b/cmd/kots/cli/install.go
@@ -46,6 +46,7 @@ import (
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 	authorizationv1 "k8s.io/api/authorization/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 )
@@ -277,6 +278,16 @@ func InstallCmd() *cobra.Command {
 				additionalAnnotations[parts[0]] = parts[1]
 			}
 
+			var tolerations []v1.Toleration
+			for _, foo := range v.GetStringSlice("tolerations") {
+				toleration, err := parseToleration(foo)
+				if err != nil {
+					return fmt.Errorf("failed to parse toleration %q: %w", foo, err)
+				}
+
+				tolerations = append(tolerations, *toleration)
+			}
+
 			deployOptions := kotsadmtypes.DeployOptions{
 				Namespace:              namespace,
 				Context:                v.GetString("context"),
@@ -308,6 +319,7 @@ func InstallCmd() *cobra.Command {
 				RequestedChannelSlug:   preferredChannelSlug,
 				AdditionalLabels:       additionalLabels,
 				AdditionalAnnotations:  additionalAnnotations,
+				Tolerations:            tolerations,
 				PrivateCAsConfigmap:    v.GetString("private-ca-configmap"),
 
 				RegistryConfig: *registryConfig,
@@ -552,6 +564,7 @@ func InstallCmd() *cobra.Command {
 	cmd.Flags().Bool("exclude-admin-console", false, "set to true to exclude the admin console and only install the application")
 	cmd.Flags().StringArray("additional-annotations", []string{}, "additional annotations to add to kotsadm pods")
 	cmd.Flags().StringArray("additional-labels", []string{}, "additional labels to add to kotsadm pods")
+	cmd.Flags().StringArray("tolerations", []string{}, "tolerations to add to kotsadm pods")
 	cmd.Flags().String("private-ca-configmap", "", "the name of a configmap containing private CAs to add to the kotsadm deployment")
 
 	registryFlags(cmd.Flags())
@@ -1110,3 +1123,69 @@ func checkPreflightResults(response *handlers.GetPreflightResultResponse, skipPr
 
 	return true, nil
 }
+
+// Parses a string in the format below into a k8s Toleration:
+// "key:operator:value:effect:tolerationSeconds" where the value can be empty and the tolerationsSeconds is optional.
+func parseToleration(input string) (*v1.Toleration, error) {
+	// Trim any leading or trailing spaces
+	input = strings.TrimSpace(input)
+
+	// Split the input string by the ':' delimiter
+	parts := strings.Split(input, ":")
+	if len(parts) < 4 {
+		return nil, fmt.Errorf("invalid toleration format, expected at least 4 fields")
+	}
+
+	operatorStr := parts[1]
+	if !isValidOperator(operatorStr) {
+		return nil, fmt.Errorf("invalid toleration operator: %q", operatorStr)
+	}
+
+	// Initialize the Toleration struct
+	toleration := &v1.Toleration{
+		Key:      parts[0],
+		Operator: v1.TolerationOperator(operatorStr),
+	}
+
+	// Handle the case where the value field might be empty
+	if len(parts) > 2 && parts[2] != "" {
+		toleration.Value = parts[2]
+	} else {
+		toleration.Value = "" // If empty, explicitly set to empty string
+	}
+
+	effectStr := parts[3]
+	if !isValidEffect(effectStr) {
+		return nil, fmt.Errorf("invalid toleration effect: %q", effectStr)
+	}
+	toleration.Effect = v1.TaintEffect(effectStr)
+
+	// If there is a fifth part, it represents tolerationSeconds
+	if len(parts) > 4 {
+		tolerationSeconds, err := strconv.ParseInt(parts[4], 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("invalid tolerationSeconds value %q: %v", parts[4], err)
+		}
+		toleration.TolerationSeconds = &tolerationSeconds
+	}
+
+	// Return the parsed toleration struct
+	return toleration, nil
+}
+
+func isValidEffect(effect string) bool {
+	validEffects := map[string]bool{
+		"NoSchedule":       true,
+		"PreferNoSchedule": true,
+		"NoExecute":        true,
+	}
+	return validEffects[effect]
+}
+
+func isValidOperator(op string) bool {
+	validOperators := map[string]bool{
+		"Exists": true,
+		"Equal":  true,
+	}
+	return validOperators[op]
+}
diff --git a/cmd/kots/cli/install_inpackage_test.go b/cmd/kots/cli/install_inpackage_test.go
@@ -0,0 +1,86 @@
+package cli
+
+import (
+	"fmt"
+	"github.com/replicatedhq/kots/pkg/util"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	v1 "k8s.io/api/core/v1"
+	"testing"
+)
+
+func Test_parseToleration(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    *v1.Toleration
+		wantErr assert.ErrorAssertionFunc
+	}{
+		{
+			name:  "Equal",
+			input: "test.com/role:Equal:core:NoSchedule",
+			want: &v1.Toleration{
+				Key:      "test.com/role",
+				Operator: v1.TolerationOpEqual,
+				Value:    "core",
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name:  "Equal 30 seconds",
+			input: "test.com/role:Equal:core:NoSchedule:30",
+			want: &v1.Toleration{
+				Key:               "test.com/role",
+				Operator:          v1.TolerationOpEqual,
+				Value:             "core",
+				Effect:            v1.TaintEffectNoSchedule,
+				TolerationSeconds: util.IntPointer(30),
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name:  "Exists",
+			input: "test.com/productid:Exists::NoSchedule",
+			want: &v1.Toleration{
+				Key:      "test.com/productid",
+				Operator: v1.TolerationOpExists,
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+		},
+		{
+			name:  "Exists with value",
+			input: "test.com/productid:Exists:testval:NoSchedule",
+			want: &v1.Toleration{
+				Key:      "test.com/productid",
+				Operator: v1.TolerationOpExists,
+				Value:    "testval",
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+		},
+		{
+			name:  "Exists 60 seconds",
+			input: "test.com/productid:Exists::NoSchedule:60",
+			want: &v1.Toleration{
+				Key:               "test.com/productid",
+				Operator:          v1.TolerationOpExists,
+				Effect:            v1.TaintEffectNoSchedule,
+				TolerationSeconds: util.IntPointer(60),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := require.New(t)
+			got, err := parseToleration(tt.input)
+			if tt.wantErr == nil {
+				req.NoError(err)
+			} else {
+				if !tt.wantErr(t, err, fmt.Sprintf("parseToleration(%v)", tt.input)) {
+					return
+				}
+			}
+			req.Equal(tt.want, got)
+		})
+	}
+}
diff --git a/pkg/kotsadm/objects/distribution_objects.go b/pkg/kotsadm/objects/distribution_objects.go
@@ -137,6 +137,7 @@ func DistributionStatefulset(deployOptions types.DeployOptions, size resource.Qu
 				},
 				Spec: corev1.PodSpec{
 					SecurityContext: securityContext,
+					Tolerations:     deployOptions.Tolerations,
 					Containers: []corev1.Container{
 						{
 							Name:            "docker-registry",

diff --git a/pkg/kotsadm/objects/kotsadm_objects.go b/pkg/kotsadm/objects/kotsadm_objects.go
@@ -450,6 +450,7 @@ func KotsadmDeployment(deployOptions types.DeployOptions) (*appsv1.Deployment, e
 					Affinity: &corev1.Affinity{
 						NodeAffinity: defaultKOTSNodeAffinity(),
 					},
+					Tolerations:        deployOptions.Tolerations,
 					SecurityContext:    securityContext,
 					Volumes:            volumes,
 					ServiceAccountName: "kotsadm",
@@ -1006,6 +1007,7 @@ func KotsadmStatefulSet(deployOptions types.DeployOptions, size resource.Quantit
 					Affinity: &corev1.Affinity{
 						NodeAffinity: defaultKOTSNodeAffinity(),
 					},
+					Tolerations:        deployOptions.Tolerations,
 					SecurityContext:    securityContext,
 					Volumes:            volumes,
 					ServiceAccountName: "kotsadm",

diff --git a/pkg/kotsadm/objects/minio_objects.go b/pkg/kotsadm/objects/minio_objects.go
@@ -143,6 +143,7 @@ func MinioStatefulset(deployOptions types.DeployOptions, size resource.Quantity)
 					Affinity: &corev1.Affinity{
 						NodeAffinity: defaultKOTSNodeAffinity(),
 					},
+					Tolerations:      deployOptions.Tolerations,
 					SecurityContext:  securityContext,
 					ImagePullSecrets: pullSecrets,
 					InitContainers:   initContainers,

diff --git a/pkg/kotsadm/objects/rqlite_objects.go b/pkg/kotsadm/objects/rqlite_objects.go
@@ -104,6 +104,7 @@ func RqliteStatefulset(deployOptions types.DeployOptions, size resource.Quantity
 					Labels:      types.GetKotsadmLabels(podLabels),
 				},
 				Spec: corev1.PodSpec{
+					Tolerations:      deployOptions.Tolerations,
 					SecurityContext:  securityContext,
 					ImagePullSecrets: pullSecrets,
 					Volumes:          volumes,

diff --git a/pkg/kotsadm/types/deployoptions.go b/pkg/kotsadm/types/deployoptions.go
@@ -59,6 +59,7 @@ type DeployOptions struct {
 	RequestedChannelSlug   string
 	AdditionalAnnotations  map[string]string
 	AdditionalLabels       map[string]string
+	Tolerations            []corev1.Toleration
 	PrivateCAsConfigmap    string
 
 	IdentityConfig kotsv1beta1.IdentityConfig