From ffbb9614d08d294b5148a58d8be8774fd39db865 Mon Sep 17 00:00:00 2001
From: Alexey Makhov <amakhov@mirantis.com>
Date: Wed, 27 Dec 2023 15:03:01 +0200
Subject: [PATCH] K0smotron pod affinities added

Signed-off-by: Alexey Makhov <amakhov@mirantis.com>
---
 .../k0smotroncluster_statefulset.go           | 22 +++++++++++++++++++
 internal/controller/k0smotron.io/util.go      |  9 +++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/internal/controller/k0smotron.io/k0smotroncluster_statefulset.go b/internal/controller/k0smotron.io/k0smotroncluster_statefulset.go
index 5e0c4116..e7df814e 100644
--- a/internal/controller/k0smotron.io/k0smotroncluster_statefulset.go
+++ b/internal/controller/k0smotron.io/k0smotroncluster_statefulset.go
@@ -77,6 +77,28 @@ func (r *ClusterReconciler) generateStatefulSet(kmc *km.Cluster) (apps.StatefulS
 					Labels: labels,
 				},
 				Spec: v1.PodSpec{
+					Affinity: &v1.Affinity{PodAntiAffinity: &v1.PodAntiAffinity{
+						PreferredDuringSchedulingIgnoredDuringExecution: []v1.WeightedPodAffinityTerm{
+							{
+								Weight: 100,
+								PodAffinityTerm: v1.PodAffinityTerm{
+									TopologyKey: "topology.kubernetes.io/zone",
+									LabelSelector: &metav1.LabelSelector{
+										MatchLabels: defaultClusterLabels(kmc),
+									},
+								},
+							},
+							{
+								Weight: 50,
+								PodAffinityTerm: v1.PodAffinityTerm{
+									TopologyKey: "kubernetes.io/hostname",
+									LabelSelector: &metav1.LabelSelector{
+										MatchLabels: defaultClusterLabels(kmc),
+									},
+								},
+							},
+						},
+					}},
 					Volumes: []v1.Volume{{
 						Name: kmc.GetEntrypointConfigMapName(),
 						VolumeSource: v1.VolumeSource{
diff --git a/internal/controller/k0smotron.io/util.go b/internal/controller/k0smotron.io/util.go
index 5369cd53..87dfee7b 100644
--- a/internal/controller/k0smotron.io/util.go
+++ b/internal/controller/k0smotron.io/util.go
@@ -18,8 +18,15 @@ package k0smotronio
 
 import km "github.com/k0sproject/k0smotron/api/k0smotron.io/v1beta1"
 
+func defaultClusterLabels(kmc *km.Cluster) map[string]string {
+	return map[string]string{
+		"app":     "k0smotron",
+		"cluster": kmc.Name,
+	}
+}
+
 func labelsForCluster(kmc *km.Cluster) map[string]string {
-	labels := map[string]string{"app": "k0smotron", "cluster": kmc.Name}
+	labels := defaultClusterLabels(kmc)
 	for k, v := range kmc.Labels {
 		labels[k] = v
 	}