added agentDaemonsetTolerationKeyEnv Var to be able to set a Toleration Key

added new option to example rook-operator.yaml

updated tests for the AGENT_TOLERATION_KEY env var

cluster/examples/kubernetes/rook-operator.yaml

@@ -124,9 +124,13 @@ spec:
         # To disable RBAC, uncomment the following:
         # - name: RBAC_ENABLED
         #  value: "false"
-        # Rook Agent toleration. Choose between NoSchedule, PreferNoSchedule and NoExecute:
+        # Rook Agent toleration. Will tolerate all taints with all keys.
+        # Choose between NoSchedule, PreferNoSchedule and NoExecute:
         # - name: AGENT_TOLERATION
         #  value: "NoSchedule"
+        # (Optional) Rook Agent toleration key. Set this to the key of the taint you want to tolerate
+        # - name: AGENT_TOLERATION_KEY
+        #  value: "<KeyOfTheTaintToTolerate"
         # The interval to check if every mon is in the quorum.
         - name: ROOK_MON_HEALTHCHECK_INTERVAL
           value: "45s"

pkg/operator/agent/agent.go

@@ -36,10 +36,11 @@ import (
 )
 
 const (
-	agentDaemonsetName          = "rook-agent"
-	flexvolumePathDirEnv        = "FLEXVOLUME_DIR_PATH"
-	flexvolumeDefaultDirPath    = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
-	agentDaemonsetTolerationEnv = "AGENT_TOLERATION"
+	agentDaemonsetName             = "rook-agent"
+	flexvolumePathDirEnv           = "FLEXVOLUME_DIR_PATH"
+	flexvolumeDefaultDirPath       = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
+	agentDaemonsetTolerationEnv    = "AGENT_TOLERATION"
+	agentDaemonsetTolerationKeyEnv = "AGENT_TOLERATION_KEY"
 )
 
 var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-agent")
@@ -187,6 +188,7 @@ func (a *Agent) createAgentDaemonSet(namespace, agentImage string) error {
 			{
 				Effect:   v1.TaintEffect(tolerationValue),
 				Operator: v1.TolerationOpExists,
+				Key:      os.Getenv(agentDaemonsetTolerationKeyEnv),
 			},
 		}
 	}

pkg/operator/agent/agent_test.go

@@ -169,6 +169,9 @@ func TestStartAgentDaemonsetWithToleration(t *testing.T) {
 	os.Setenv(agentDaemonsetTolerationEnv, "NoSchedule")
 	defer os.Unsetenv(agentDaemonsetTolerationEnv)
 
+	os.Setenv(agentDaemonsetTolerationKeyEnv, "example")
+	defer os.Unsetenv(agentDaemonsetTolerationKeyEnv)
+
 	pod := v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      "rook-operator",
@@ -197,5 +200,6 @@ func TestStartAgentDaemonsetWithToleration(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 1, len(agentDS.Spec.Template.Spec.Tolerations))
 	assert.Equal(t, "NoSchedule", string(agentDS.Spec.Template.Spec.Tolerations[0].Effect))
+	assert.Equal(t, "example", string(agentDS.Spec.Template.Spec.Tolerations[0].Key))
 	assert.Equal(t, "Exists", string(agentDS.Spec.Template.Spec.Tolerations[0].Operator))
 }