Add node affinity for members pod (#494)

Author: informalict

pkg/util/k8sutil/affinity.go

@@ -23,7 +23,7 @@
 package k8sutil
 
 import (
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -32,6 +32,21 @@ import (
 // affinityWithRole contains the role to configure affinity with.
 func createAffinity(deploymentName, role string, required bool, affinityWithRole string) *v1.Affinity {
 	a := &v1.Affinity{
+		NodeAffinity: &v1.NodeAffinity{
+			RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
+				NodeSelectorTerms: []v1.NodeSelectorTerm{
+					{
+						MatchExpressions: []v1.NodeSelectorRequirement{
+							{
+								Key:      "beta.kubernetes.io/arch",
+								Operator: "In",
+								Values:   []string{"amd64"},
+							},
+						},
+					},
+				},
+			},
+		},
 		PodAntiAffinity: &v1.PodAntiAffinity{},
 	}
 	labels := LabelsForDeployment(deploymentName, role)
@@ -72,5 +87,6 @@ func createAffinity(deploymentName, role string, required bool, affinityWithRole
 			})
 		}
 	}
+
 	return a
 }

pkg/util/k8sutil/affinity_test.go

@@ -25,12 +25,29 @@ package k8sutil
 import (
 	"testing"
 
+	v1 "k8s.io/api/core/v1"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 // TestCreateAffinity tests createAffinity
 func TestCreateAffinity(t *testing.T) {
+	expectedNodeAffinity := &v1.NodeAffinity{
+		RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
+			NodeSelectorTerms: []v1.NodeSelectorTerm{
+				{
+					MatchExpressions: []v1.NodeSelectorRequirement{
+						{
+							Key:      "beta.kubernetes.io/arch",
+							Operator: "In",
+							Values:   []string{"amd64"},
+						},
+					},
+				},
+			},
+		},
+	}
 	// Required
 	a := createAffinity("test", "role", true, "")
 	assert.Nil(t, a.PodAffinity)
@@ -42,6 +59,7 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "role", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 
 	// Require & affinity with role dbserver
 	a = createAffinity("test", "role", true, "dbserver")
@@ -53,6 +71,7 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "dbserver", a.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 
 	require.NotNil(t, a.PodAntiAffinity)
 	require.Len(t, a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution, 1)
@@ -62,6 +81,7 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "role", a.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution[0].LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 
 	// Not Required
 	a = createAffinity("test", "role", false, "")
@@ -74,6 +94,7 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "role", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 
 	// Not Required & affinity with role dbserver
 	a = createAffinity("test", "role", false, "dbserver")
@@ -85,6 +106,7 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "dbserver", a.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 
 	require.NotNil(t, a.PodAntiAffinity)
 	require.Len(t, a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 1)
@@ -94,4 +116,5 @@ func TestCreateAffinity(t *testing.T) {
 	assert.Equal(t, "test", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["arango_deployment"])
 	assert.Equal(t, "arangodb", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["app"])
 	assert.Equal(t, "role", a.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[0].PodAffinityTerm.LabelSelector.MatchLabels["role"])
+	assert.Equal(t, expectedNodeAffinity, a.NodeAffinity)
 }