From b73d5ddeb5ca4755ab769fc98478f465d2dddc6f Mon Sep 17 00:00:00 2001
From: Steve Leon <kokhang@gmail.com>
Date: Tue, 12 Dec 2017 13:30:29 -0800
Subject: [PATCH] Add toleration for agent daemonset

Fixes https://github.com/rook/rook/issues/1295
---
 cluster/charts/rook/templates/deployment.yaml |  6 +++
 cluster/charts/rook/values.yaml.tmpl          |  5 +++
 .../examples/kubernetes/rook-operator.yaml    |  3 ++
 pkg/operator/agent/agent.go                   | 18 ++++++--
 pkg/operator/agent/agent_test.go              | 44 +++++++++++++++++++
 5 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/cluster/charts/rook/templates/deployment.yaml b/cluster/charts/rook/templates/deployment.yaml
index fb542a015..6c77f9421 100644
--- a/cluster/charts/rook/templates/deployment.yaml
+++ b/cluster/charts/rook/templates/deployment.yaml
@@ -26,6 +26,12 @@ spec:
 {{- if not .Values.rbacEnable }}
         - name: RBAC_ENABLED
           value: "false"
+{{- end }}
+{{- if .Values.agent }}
+{{- if .Values.agent.toleration }}
+        - name: AGENT_TOLERATION
+          value: {{ .Values.agent.toleration }}
+{{- end }}
 {{- end }}
         - name: ROOK_LOG_LEVEL
           value: {{ .Values.logLevel }}
diff --git a/cluster/charts/rook/values.yaml.tmpl b/cluster/charts/rook/values.yaml.tmpl
index a5fa9622d..d0d4d584d 100644
--- a/cluster/charts/rook/values.yaml.tmpl
+++ b/cluster/charts/rook/values.yaml.tmpl
@@ -26,3 +26,8 @@ logLevel: INFO
 ## If true, create & use RBAC resources
 ##
 rbacEnable: true
+
+## Rook Agent configuration
+## toleration. Choose between NoSchedule, PreferNoSchedule and NoExecute:
+# agent:
+#   toleration: NoSchedule
diff --git a/cluster/examples/kubernetes/rook-operator.yaml b/cluster/examples/kubernetes/rook-operator.yaml
index a80af838a..5ede0b1f8 100644
--- a/cluster/examples/kubernetes/rook-operator.yaml
+++ b/cluster/examples/kubernetes/rook-operator.yaml
@@ -126,6 +126,9 @@ spec:
         # To disable RBAC, uncomment the following:
         # - name: RBAC_ENABLED
         #  value: "false"
+        # Rook Agent toleration. Choose between NoSchedule, PreferNoSchedule and NoExecute:
+        # - name: AGENT_TOLERATION
+        #  value: "NoSchedule"
         # The interval to check if every mon is in the quorum.
         - name: ROOK_MON_HEALTHCHECK_INTERVAL
           value: "45s"
diff --git a/pkg/operator/agent/agent.go b/pkg/operator/agent/agent.go
index 41d5b09e5..84343f58d 100644
--- a/pkg/operator/agent/agent.go
+++ b/pkg/operator/agent/agent.go
@@ -38,9 +38,10 @@ import (
 )
 
 const (
-	agentDaemonsetName       = "rook-agent"
-	flexvolumePathDirEnv     = "FLEXVOLUME_DIR_PATH"
-	flexvolumeDefaultDirPath = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
+	agentDaemonsetName          = "rook-agent"
+	flexvolumePathDirEnv        = "FLEXVOLUME_DIR_PATH"
+	flexvolumeDefaultDirPath    = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
+	agentDaemonsetTolerationEnv = "AGENT_TOLERATION"
 )
 
 var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-agent")
@@ -187,6 +188,17 @@ func (a *Agent) createAgentDaemonSet(namespace string) error {
 		},
 	}
 
+	// Add toleration if any
+	tolerationValue := os.Getenv(agentDaemonsetTolerationEnv)
+	if tolerationValue != "" {
+		ds.Spec.Template.Spec.Tolerations = []v1.Toleration{
+			{
+				Effect:   v1.TaintEffect(tolerationValue),
+				Operator: v1.TolerationOpExists,
+			},
+		}
+	}
+
 	_, err = a.clientset.Extensions().DaemonSets(namespace).Create(ds)
 	if err != nil {
 		if !kserrors.IsAlreadyExists(err) {
diff --git a/pkg/operator/agent/agent_test.go b/pkg/operator/agent/agent_test.go
index 9ab6e8e29..0b65316f5 100644
--- a/pkg/operator/agent/agent_test.go
+++ b/pkg/operator/agent/agent_test.go
@@ -91,6 +91,7 @@ func TestStartAgentDaemonset(t *testing.T) {
 	assert.Equal(t, 2, len(envs))
 	image := agentDS.Spec.Template.Spec.Containers[0].Image
 	assert.Equal(t, "rook/test", image)
+	assert.Nil(t, agentDS.Spec.Template.Spec.Tolerations)
 }
 
 func TestGetContainerImage(t *testing.T) {
@@ -158,3 +159,46 @@ func TestGetContainerImageMultipleContainers(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.Equal(t, "failed to get container image. There should only be exactly one container in this pod", err.Error())
 }
+
+func TestStartAgentDaemonsetWithToleration(t *testing.T) {
+	clientset := test.New(3)
+
+	os.Setenv(k8sutil.PodNamespaceEnvVar, "rook-system")
+	defer os.Unsetenv(k8sutil.PodNamespaceEnvVar)
+
+	os.Setenv(k8sutil.PodNameEnvVar, "rook-operator")
+	defer os.Unsetenv(k8sutil.PodNameEnvVar)
+
+	os.Setenv(agentDaemonsetTolerationEnv, "NoSchedule")
+	defer os.Unsetenv(agentDaemonsetTolerationEnv)
+
+	pod := v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "rook-operator",
+			Namespace: "rook-system",
+		},
+		Spec: v1.PodSpec{
+			Containers: []v1.Container{
+				{
+					Name:  "mypodContainer",
+					Image: "rook/test",
+				},
+			},
+		},
+	}
+	clientset.CoreV1().Pods("rook-system").Create(&pod)
+
+	namespace := "ns"
+	a := New(clientset)
+
+	// start a basic cluster
+	err := a.Start(namespace)
+	assert.Nil(t, err)
+
+	// check daemonset toleration
+	agentDS, err := clientset.Extensions().DaemonSets(namespace).Get("rook-agent", metav1.GetOptions{})
+	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, "Exists", string(agentDS.Spec.Template.Spec.Tolerations[0].Operator))
+}