feat: add pod-anti-affinity

openGemini · Sep 19, 2023 · 570dbee · 570dbee
1 parent 6555659
commit 570dbee
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 3 deletions.
diff --git a/api/v1/geminicluster_types.go b/api/v1/geminicluster_types.go
@@ -214,6 +214,10 @@ func (cluster *GeminiCluster) GetEnableHttpAuth() bool {
 	return false
 }
 
+func (cluster *GeminiCluster) GetEnableAffinity() bool {
+	return cluster.Spec.Affinity.EnablePodAntiAffinity
+}
+
 func (cluster *GeminiCluster) GetAdminUserSecretName() string {
 	return fmt.Sprintf("%v%v", cluster.Name, AdminUserSecretSuffix)
 }

diff --git a/controllers/instance.go b/controllers/instance.go
@@ -3,15 +3,17 @@ package controllers
 import (
 	"context"
 
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
 	opengeminiv1 "github.com/openGemini/openGemini-operator/api/v1"
 	"github.com/openGemini/openGemini-operator/pkg/naming"
 	"github.com/openGemini/openGemini-operator/pkg/opengemini/meta"
 	"github.com/openGemini/openGemini-operator/pkg/opengemini/sql"
 	"github.com/openGemini/openGemini-operator/pkg/opengemini/store"
+	"github.com/openGemini/openGemini-operator/pkg/specs"
 	"github.com/openGemini/openGemini-operator/pkg/utils"
-	appsv1 "k8s.io/api/apps/v1"
-	corev1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 // +kubebuilder:rbac:groups="apps",resources="statefulsets",verbs={get,create,patch}
@@ -137,6 +139,7 @@ func generateInstanceStatefulSetIntent(
 			opengeminiv1.LabelInstance:    sts.Name,
 			opengeminiv1.LabelConfigHash:  cluster.Status.AppliedConfigHash,
 		})
+	sts.Spec.Template.Spec.Affinity = specs.CreateAffinity(cluster.GetEnableAffinity(), sts.Spec.Template.Labels)
 	sts.Spec.Template.Spec.RestartPolicy = corev1.RestartPolicyAlways
 	sts.Spec.Template.Spec.ShareProcessNamespace = &[]bool{true}[0]
 	sts.Spec.Template.Spec.EnableServiceLinks = &[]bool{false}[0]
@@ -178,6 +181,7 @@ func generateInstanceDeploymentIntent(
 			opengeminiv1.LabelInstance:    deploy.Name,
 			opengeminiv1.LabelConfigHash:  cluster.Status.AppliedConfigHash,
 		})
+	deploy.Spec.Template.Spec.Affinity = specs.CreateAffinity(cluster.GetEnableAffinity(), deploy.Spec.Template.Labels)
 	deploy.Spec.Template.Spec.RestartPolicy = corev1.RestartPolicyAlways
 	deploy.Spec.Template.Spec.ShareProcessNamespace = &[]bool{true}[0]
 	deploy.Spec.Template.Spec.EnableServiceLinks = &[]bool{false}[0]

diff --git a/pkg/specs/affinity.go b/pkg/specs/affinity.go
@@ -0,0 +1,35 @@
+package specs
+
+import (
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+const defaultAffinityWeight = 100
+
+func CreateAffinity(enable bool, matchLabels map[string]string) *corev1.Affinity {
+	affinity := corev1.Affinity{
+		PodAntiAffinity: &corev1.PodAntiAffinity{},
+	}
+
+	affinityTerm := corev1.PodAffinityTerm{
+		TopologyKey: corev1.LabelHostname,
+		LabelSelector: &metav1.LabelSelector{
+			MatchLabels: matchLabels,
+		},
+	}
+
+	if enable {
+		affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = []corev1.PodAffinityTerm{
+			affinityTerm,
+		}
+	} else {
+		affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = []corev1.WeightedPodAffinityTerm{
+			{
+				Weight:          defaultAffinityWeight,
+				PodAffinityTerm: affinityTerm,
+			},
+		}
+	}
+	return &affinity
+}