feat: add pod-anti-affinity (#21)

Signed-off-by: Xie Wei <[email protected]>
openGemini · Oct 12, 2023 · c8059da · c8059da
1 parent d338bc6
commit c8059da
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 28 deletions.
diff --git a/api/v1/geminicluster_types.go b/api/v1/geminicluster_types.go
@@ -215,6 +215,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
@@ -12,6 +12,7 @@ import (
 	"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"
 )
 
@@ -113,31 +114,35 @@ func generateInstanceStatefulSetIntent(
 ) {
 	sts.Annotations = utils.MergeLabels(
 		cluster.Spec.Metadata.GetAnnotationsOrNil())
+
+	baseLabels := map[string]string{
+		opengeminiv1.LabelCluster:     cluster.Name,
+		opengeminiv1.LabelInstanceSet: setName,
+	}
+	matchLabels := utils.MergeLabels(
+		baseLabels,
+		map[string]string{
+			opengeminiv1.LabelInstance: sts.Name,
+		},
+	)
+
 	sts.Labels = utils.MergeLabels(
 		cluster.Spec.Metadata.GetLabelsOrNil(),
-		map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    sts.Name,
-		})
+		matchLabels,
+	)
 	sts.Spec.Selector = &metav1.LabelSelector{
-		MatchLabels: map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    sts.Name,
-		},
+		MatchLabels: matchLabels,
 	}
 	sts.Spec.Template.Annotations = utils.MergeLabels(
 		cluster.Spec.Metadata.GetAnnotationsOrNil(),
 	)
 	sts.Spec.Template.Labels = utils.MergeLabels(
 		cluster.Spec.Metadata.GetLabelsOrNil(),
+		matchLabels,
 		map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    sts.Name,
-			opengeminiv1.LabelConfigHash:  cluster.Status.AppliedConfigHash,
+			opengeminiv1.LabelConfigHash: cluster.Status.AppliedConfigHash,
 		})
+	sts.Spec.Template.Spec.Affinity = specs.CreateAffinity(cluster.GetEnableAffinity(), baseLabels)
 	sts.Spec.Template.Spec.RestartPolicy = corev1.RestartPolicyAlways
 	sts.Spec.Template.Spec.ShareProcessNamespace = &[]bool{true}[0]
 	sts.Spec.Template.Spec.EnableServiceLinks = &[]bool{false}[0]
@@ -154,31 +159,35 @@ func generateInstanceDeploymentIntent(
 ) {
 	deploy.Annotations = utils.MergeLabels(
 		cluster.Spec.Metadata.GetAnnotationsOrNil())
+
+	baseLabels := map[string]string{
+		opengeminiv1.LabelCluster:     cluster.Name,
+		opengeminiv1.LabelInstanceSet: setName,
+	}
+	matchLabels := utils.MergeLabels(
+		baseLabels,
+		map[string]string{
+			opengeminiv1.LabelInstance: deploy.Name,
+		},
+	)
+
 	deploy.Labels = utils.MergeLabels(
 		cluster.Spec.Metadata.GetLabelsOrNil(),
-		map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    deploy.Name,
-		})
+		matchLabels,
+	)
 	deploy.Spec.Selector = &metav1.LabelSelector{
-		MatchLabels: map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    deploy.Name,
-		},
+		MatchLabels: matchLabels,
 	}
 	deploy.Spec.Template.Annotations = utils.MergeLabels(
 		cluster.Spec.Metadata.GetAnnotationsOrNil(),
 	)
 	deploy.Spec.Template.Labels = utils.MergeLabels(
 		cluster.Spec.Metadata.GetLabelsOrNil(),
+		matchLabels,
 		map[string]string{
-			opengeminiv1.LabelCluster:     cluster.Name,
-			opengeminiv1.LabelInstanceSet: setName,
-			opengeminiv1.LabelInstance:    deploy.Name,
-			opengeminiv1.LabelConfigHash:  cluster.Status.AppliedConfigHash,
+			opengeminiv1.LabelConfigHash: cluster.Status.AppliedConfigHash,
 		})
+	deploy.Spec.Template.Spec.Affinity = specs.CreateAffinity(cluster.GetEnableAffinity(), baseLabels)
 	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
+}