feat: add default affinity (#33)

zncdatadev · May 29, 2024 · 23ae2b2 · 23ae2b2
1 parent a5977ae
commit 23ae2b2
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 4 deletions.
diff --git a/internal/common/labels.go b/internal/common/labels.go
@@ -4,19 +4,25 @@ import (
 	"strings"
 )
 
+const (
+	LabelCrName    = "app.kubernetes.io/Name"
+	LabelComponent = "app.kubernetes.io/component"
+	LabelManagedBy = "app.kubernetes.io/managed-by"
+)
+
 type RoleLabels struct {
 	InstanceName string
 	Name         string
 }
 
 func (r *RoleLabels) GetLabels() map[string]string {
 	res := map[string]string{
-		"app.kubernetes.io/Name":       strings.ToLower(r.InstanceName),
-		"app.kubernetes.io/component":  r.Name,
-		"app.kubernetes.io/managed-by": "hdfs-operator",
+		LabelCrName:    strings.ToLower(r.InstanceName),
+		LabelComponent: r.Name,
+		LabelManagedBy: "hdfs-operator",
 	}
 	if r.Name != "" {
-		res["app.kubernetes.io/component"] = r.Name
+		res[LabelComponent] = r.Name
 	}
 	return res
 }

diff --git a/internal/common/resouce_reconciler.go b/internal/common/resouce_reconciler.go
@@ -57,6 +57,10 @@ type ConditionsGetter interface {
 	GetConditions() *[]metav1.Condition
 }
 
+type AffinitySetter interface {
+	SetAffinity(resource client.Object)
+}
+
 type WorkloadOverride interface {
 	CommandOverride(resource client.Object)
 	EnvOverride(resource client.Object)
@@ -267,6 +271,9 @@ func (s *WorkloadStyleReconciler[T, G]) DoReconcile(
 	// check if the resource is satisfied
 	// if not, return requeue
 	// if satisfied, return nil
+	if setAffinity, ok := instance.(AffinitySetter); ok {
+		setAffinity.SetAffinity(resource)
+	}
 	if override, ok := instance.(WorkloadOverride); ok {
 		override.CommandOverride(resource)
 		override.EnvOverride(resource)

diff --git a/internal/common/util.go b/internal/common/util.go
@@ -6,6 +6,7 @@ import (
 	"github.com/zncdata-labs/hdfs-operator/internal/util"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"strings"
 )
 
@@ -258,3 +259,39 @@ func NameNodePodNames(instanceName string, groupName string) []string {
 	pods := CreatePodNamesByReplicas(naneNodeReplicas, nameNodeStatefulSetName)
 	return pods
 }
+
+func AffinityDefault(role Role, crName string) *corev1.Affinity {
+	return &corev1.Affinity{
+		PodAffinity: &corev1.PodAffinity{
+			PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{
+				{
+					Weight: 20,
+					PodAffinityTerm: corev1.PodAffinityTerm{
+						LabelSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								LabelCrName: crName,
+							},
+						},
+						TopologyKey: corev1.LabelHostname,
+					},
+				},
+			},
+		},
+		PodAntiAffinity: &corev1.PodAntiAffinity{
+			PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{
+				{
+					Weight: 70,
+					PodAffinityTerm: corev1.PodAffinityTerm{
+						LabelSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								LabelCrName:    crName,
+								LabelComponent: string(role),
+							},
+						},
+						TopologyKey: corev1.LabelHostname,
+					},
+				},
+			},
+		},
+	}
+}
diff --git a/internal/controller/data/statefulset.go b/internal/controller/data/statefulset.go
@@ -72,6 +72,16 @@ func (s *StatefulSetReconciler) Build(_ context.Context) (client.Object, error)
 		},
 	}, nil
 }
+
+func (s *StatefulSetReconciler) SetAffinity(resource client.Object) {
+	dep := resource.(*appv1.StatefulSet)
+	if affinity := s.MergedCfg.Config.Affinity; affinity != nil {
+		dep.Spec.Template.Spec.Affinity = affinity
+	} else {
+		dep.Spec.Template.Spec.Affinity = common.AffinityDefault(common.DataNode, s.Instance.GetName())
+	}
+}
+
 func (s *StatefulSetReconciler) CommandOverride(resource client.Object) {
 	dep := resource.(*appv1.StatefulSet)
 	containers := dep.Spec.Template.Spec.Containers

diff --git a/internal/controller/journal/statefulset.go b/internal/controller/journal/statefulset.go
@@ -69,6 +69,16 @@ func (s *StatefulSetReconciler) Build(_ context.Context) (client.Object, error)
 		},
 	}, nil
 }
+
+func (s *StatefulSetReconciler) SetAffinity(resource client.Object) {
+	dep := resource.(*appv1.StatefulSet)
+	if affinity := s.MergedCfg.Config.Affinity; affinity != nil {
+		dep.Spec.Template.Spec.Affinity = affinity
+	} else {
+		dep.Spec.Template.Spec.Affinity = common.AffinityDefault(common.JournalNode, s.Instance.GetName())
+	}
+}
+
 func (s *StatefulSetReconciler) CommandOverride(resource client.Object) {
 	dep := resource.(*appv1.StatefulSet)
 	containers := dep.Spec.Template.Spec.Containers

diff --git a/internal/controller/name/statefulset.go b/internal/controller/name/statefulset.go
@@ -79,6 +79,16 @@ func (s *StatefulSetReconciler) Build(_ context.Context) (client.Object, error)
 		},
 	}, nil
 }
+
+func (s *StatefulSetReconciler) SetAffinity(resource client.Object) {
+	dep := resource.(*appv1.StatefulSet)
+	if affinity := s.MergedCfg.Config.Affinity; affinity != nil {
+		dep.Spec.Template.Spec.Affinity = affinity
+	} else {
+		dep.Spec.Template.Spec.Affinity = common.AffinityDefault(common.NameNode, s.Instance.GetName())
+	}
+}
+
 func (s *StatefulSetReconciler) CommandOverride(resource client.Object) {
 	dep := resource.(*appv1.StatefulSet)
 	containers := dep.Spec.Template.Spec.Containers