operator: Add label selector to zone awareness TopologySpreadConstraints (#9406)

operator/CHANGELOG.md

@@ -1,5 +1,6 @@
 ## Main
 
+- [9406](https://github.com/grafana/loki/pull/9406) **aminesnow**: Add label selector to zone awareness TopologySpreadConstraints
 - [9366](https://github.com/grafana/loki/pull/9366) **periklis**: Add support for custom tenant topology in rules
 - [9315](https://github.com/grafana/loki/pull/9315) **aminesnow**: Add zone awareness spec to LokiStack
 - [9343](https://github.com/grafana/loki/pull/9343) **JoaoBraveCoding**: Add default PodAntiAffinity to Query Frontend

operator/internal/manifests/distributor.go

@@ -128,7 +128,7 @@ func NewDistributorDeployment(opts Options) *appsv1.Deployment {
 	}
 
 	if opts.Stack.Replication != nil {
-		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication)
+		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication, LabelDistributorComponent, opts.Name)
 	}
 
 	return &appsv1.Deployment{

operator/internal/manifests/distributor_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 	"github.com/grafana/loki/operator/internal/manifests"
@@ -131,11 +132,23 @@ func TestNewDistributorDeployment_TopologySpreadConstraints(t *testing.T) {
 			MaxSkew:           3,
 			TopologyKey:       "zone",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "distributor",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 		{
 			MaxSkew:           2,
 			TopologyKey:       "region",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "distributor",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 	}, depl.Spec.Template.Spec.TopologySpreadConstraints)
 }

operator/internal/manifests/indexgateway.go

@@ -134,7 +134,7 @@ func NewIndexGatewayStatefulSet(opts Options) *appsv1.StatefulSet {
 	}
 
 	if opts.Stack.Replication != nil {
-		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication)
+		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication, LabelIndexGatewayComponent, opts.Name)
 	}
 
 	return &appsv1.StatefulSet{

operator/internal/manifests/indexgateway_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 	"github.com/grafana/loki/operator/internal/manifests"
@@ -137,11 +138,23 @@ func TestNewIndexGatewayStatefulSet_TopologySpreadConstraints(t *testing.T) {
 			MaxSkew:           3,
 			TopologyKey:       "zone",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "index-gateway",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 		{
 			MaxSkew:           2,
 			TopologyKey:       "region",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "index-gateway",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 	}, depl.Spec.Template.Spec.TopologySpreadConstraints)
 }

operator/internal/manifests/ingester.go

@@ -144,7 +144,7 @@ func NewIngesterStatefulSet(opts Options) *appsv1.StatefulSet {
 	}
 
 	if opts.Stack.Replication != nil {
-		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication)
+		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication, LabelIngesterComponent, opts.Name)
 	}
 
 	return &appsv1.StatefulSet{

operator/internal/manifests/ingester_test.go

@@ -167,11 +167,23 @@ func TestNewIngesterStatefulSet_TopologySpreadConstraints(t *testing.T) {
 			MaxSkew:           2,
 			TopologyKey:       "zone",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "ingester",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 		{
 			MaxSkew:           1,
 			TopologyKey:       "region",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "ingester",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 	}, ss.Spec.Template.Spec.TopologySpreadConstraints)
 }

operator/internal/manifests/querier.go

@@ -134,7 +134,7 @@ func NewQuerierDeployment(opts Options) *appsv1.Deployment {
 	}
 
 	if opts.Stack.Replication != nil {
-		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication)
+		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication, LabelQuerierComponent, opts.Name)
 	}
 
 	return &appsv1.Deployment{

operator/internal/manifests/querier_test.go

@@ -217,11 +217,23 @@ func TestNewQuerierDeployment_TopologySpreadConstraints(t *testing.T) {
 			MaxSkew:           2,
 			TopologyKey:       "zone",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "querier",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 		{
 			MaxSkew:           1,
 			TopologyKey:       "region",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "querier",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 	}, depl.Spec.Template.Spec.TopologySpreadConstraints)
 }

operator/internal/manifests/query-frontend.go

@@ -140,7 +140,7 @@ func NewQueryFrontendDeployment(opts Options) *appsv1.Deployment {
 	}
 
 	if opts.Stack.Replication != nil {
-		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication)
+		podSpec.TopologySpreadConstraints = topologySpreadConstraints(*opts.Stack.Replication, LabelQueryFrontendComponent, opts.Name)
 	}
 
 	return &appsv1.Deployment{

operator/internal/manifests/query-frontend_test.go

@@ -126,11 +126,23 @@ func TestNewQueryFrontendDeployment_TopologySpreadConstraints(t *testing.T) {
 			MaxSkew:           1,
 			TopologyKey:       "zone",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "query-frontend",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 		{
 			MaxSkew:           2,
 			TopologyKey:       "region",
 			WhenUnsatisfiable: "DoNotSchedule",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app.kubernetes.io/component": "query-frontend",
+					"app.kubernetes.io/instance":  "abcd",
+				},
+			},
 		},
 	}, depl.Spec.Template.Spec.TopologySpreadConstraints)
 }
@@ -165,4 +177,4 @@ func TestQueryFrontendPodAntiAffinity(t *testing.T) {
 	}
 	require.NotNil(t, sts.Spec.Template.Spec.Affinity)
 	require.Equal(t, expectedPodAntiAffinity, sts.Spec.Template.Spec.Affinity.PodAntiAffinity)
-}
+}

operator/internal/manifests/var.go

@@ -143,7 +143,7 @@ func serviceAnnotations(serviceName string, enableSigningService bool) map[strin
 	return annotations
 }
 
-func topologySpreadConstraints(spec lokiv1.ReplicationSpec) []corev1.TopologySpreadConstraint {
+func topologySpreadConstraints(spec lokiv1.ReplicationSpec, component string, stackName string) []corev1.TopologySpreadConstraint {
 	var tsc []corev1.TopologySpreadConstraint
 	if len(spec.Zones) > 0 {
 		tsc = make([]corev1.TopologySpreadConstraint, len(spec.Zones))
@@ -152,6 +152,12 @@ func topologySpreadConstraints(spec lokiv1.ReplicationSpec) []corev1.TopologySpr
 				MaxSkew:           int32(z.MaxSkew),
 				TopologyKey:       z.TopologyKey,
 				WhenUnsatisfiable: corev1.DoNotSchedule,
+				LabelSelector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						kubernetesCompomentLabel: component,
+						kubernetesInstanceLabel:  stackName,
+					},
+				},
 			}
 		}
 	}