operator: Add PodDisruptionBudget to lokistack-ingestion path (#9260)

operator/CHANGELOG.md

@@ -1,5 +1,6 @@
 ## Main
 - [9262](https://github.com/grafana/loki/pull/9262) **btaani**: Add PodDisruptionBudget to the Ruler
+- [9260](https://github.com/grafana/loki/pull/9260) **JoaoBraveCoding**: Add PodDisruptionBudgets to the ingestion path
 - [9188](https://github.com/grafana/loki/pull/9188) **aminesnow**: Add PodDisruptionBudgets to the query path
 - [9162](https://github.com/grafana/loki/pull/9162) **aminesnow**: Add a PodDisruptionBudget to lokistack-gateway
 - [9049](https://github.com/grafana/loki/pull/9049) **alanconway**: Revert 1x.extra-small changes, add 1x.demo

operator/internal/manifests/distributor.go

@@ -8,6 +8,7 @@ import (
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -49,6 +50,7 @@ func BuildDistributor(opts Options) ([]client.Object, error) {
 		deployment,
 		NewDistributorGRPCService(opts),
 		NewDistributorHTTPService(opts),
+		newDistributorPodDisruptionBudget(opts),
 	}, nil
 }
 
@@ -221,3 +223,27 @@ func configureDistributorGRPCServicePKI(deployment *appsv1.Deployment, opts Opti
 	serviceName := serviceNameDistributorGRPC(opts.Name)
 	return configureGRPCServicePKI(&deployment.Spec.Template.Spec, serviceName)
 }
+
+// newDistributorPodDisruptionBudget returns a PodDisruptionBudget for the LokiStack
+// Distributor pods.
+func newDistributorPodDisruptionBudget(opts Options) *policyv1.PodDisruptionBudget {
+	l := ComponentLabels(LabelDistributorComponent, opts.Name)
+	mu := intstr.FromInt(1)
+	return &policyv1.PodDisruptionBudget{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       "PodDisruptionBudget",
+			APIVersion: policyv1.SchemeGroupVersion.String(),
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Labels:    l,
+			Name:      DistributorName(opts.Name),
+			Namespace: opts.Namespace,
+		},
+		Spec: policyv1.PodDisruptionBudgetSpec{
+			Selector: &metav1.LabelSelector{
+				MatchLabels: l,
+			},
+			MinAvailable: &mu,
+		},
+	}
+}

operator/internal/manifests/distributor_test.go

@@ -1,11 +1,14 @@
 package manifests_test
 
 import (
+	"math/rand"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+	policyv1 "k8s.io/api/policy/v1"
+
 	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 	"github.com/grafana/loki/operator/internal/manifests"
-	"github.com/stretchr/testify/require"
 )
 
 func TestNewDistributorDeployment_SelectorMatchesLabels(t *testing.T) {
@@ -67,3 +70,31 @@ func TestNewDistributorDeployment_HasTemplateCertRotationRequiredAtAnnotation(t
 	require.Contains(t, annotations, expected)
 	require.Equal(t, annotations[expected], "deadbeef")
 }
+
+func TestBuildDistributor_PodDisruptionBudget(t *testing.T) {
+	opts := manifests.Options{
+		Name:      "abcd",
+		Namespace: "efgh",
+		Stack: lokiv1.LokiStackSpec{
+			Template: &lokiv1.LokiTemplateSpec{
+				Distributor: &lokiv1.LokiComponentSpec{
+					Replicas: rand.Int31(),
+				},
+			},
+			Tenants: &lokiv1.TenantsSpec{
+				Mode: lokiv1.OpenshiftLogging,
+			},
+		},
+	}
+	objs, err := manifests.BuildDistributor(opts)
+	require.NoError(t, err)
+	require.Len(t, objs, 4)
+
+	pdb := objs[3].(*policyv1.PodDisruptionBudget)
+	require.NotNil(t, pdb)
+	require.Equal(t, "abcd-distributor", pdb.Name)
+	require.Equal(t, "efgh", pdb.Namespace)
+	require.NotNil(t, pdb.Spec.MinAvailable.IntVal)
+	require.Equal(t, int32(1), pdb.Spec.MinAvailable.IntVal)
+	require.EqualValues(t, manifests.ComponentLabels(manifests.LabelDistributorComponent, opts.Name), pdb.Spec.Selector.MatchLabels)
+}

operator/internal/manifests/ingester.go

@@ -9,6 +9,7 @@ import (
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
@@ -55,6 +56,7 @@ func BuildIngester(opts Options) ([]client.Object, error) {
 		statefulSet,
 		NewIngesterGRPCService(opts),
 		NewIngesterHTTPService(opts),
+		newIngesterPodDisruptionBudget(opts),
 	}, nil
 }
 
@@ -274,3 +276,31 @@ func configureIngesterGRPCServicePKI(sts *appsv1.StatefulSet, opts Options) erro
 	serviceName := serviceNameIngesterGRPC(opts.Name)
 	return configureGRPCServicePKI(&sts.Spec.Template.Spec, serviceName)
 }
+
+// newIngesterPodDisruptionBudget returns a PodDisruptionBudget for the LokiStack
+// Ingester pods.
+func newIngesterPodDisruptionBudget(opts Options) *policyv1.PodDisruptionBudget {
+	l := ComponentLabels(LabelIngesterComponent, opts.Name)
+	// Default to 1 if not defined in ResourceRequirementsTable for a given size
+	mu := intstr.FromInt(1)
+	if opts.ResourceRequirements.Ingester.PDBMinAvailable > 0 {
+		mu = intstr.FromInt(opts.ResourceRequirements.Ingester.PDBMinAvailable)
+	}
+	return &policyv1.PodDisruptionBudget{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       "PodDisruptionBudget",
+			APIVersion: policyv1.SchemeGroupVersion.String(),
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Labels:    l,
+			Name:      IngesterName(opts.Name),
+			Namespace: opts.Namespace,
+		},
+		Spec: policyv1.PodDisruptionBudgetSpec{
+			Selector: &metav1.LabelSelector{
+				MatchLabels: l,
+			},
+			MinAvailable: &mu,
+		},
+	}
+}

operator/internal/manifests/ingester_test.go

@@ -1,11 +1,16 @@
 package manifests_test
 
 import (
+	"math/rand"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+	policyv1 "k8s.io/api/policy/v1"
+
+	v1 "github.com/grafana/loki/operator/apis/config/v1"
 	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 	"github.com/grafana/loki/operator/internal/manifests"
-	"github.com/stretchr/testify/require"
+	"github.com/grafana/loki/operator/internal/manifests/internal"
 )
 
 func TestNewIngesterStatefulSet_HasTemplateConfigHashAnnotation(t *testing.T) {
@@ -75,3 +80,56 @@ func TestNewIngesterStatefulSet_SelectorMatchesLabels(t *testing.T) {
 		require.Equal(t, l[key], value)
 	}
 }
+
+func TestBuildIngester_PodDisruptionBudget(t *testing.T) {
+	for _, tc := range []struct {
+		Name                 string
+		PDBMinAvailable      int
+		ExpectedMinAvailable int
+	}{
+		{
+			Name:                 "Small stack",
+			PDBMinAvailable:      1,
+			ExpectedMinAvailable: 1,
+		},
+		{
+			Name:                 "Medium stack",
+			PDBMinAvailable:      2,
+			ExpectedMinAvailable: 2,
+		},
+	} {
+		t.Run(tc.Name, func(t *testing.T) {
+			opts := manifests.Options{
+				Name:      "abcd",
+				Namespace: "efgh",
+				Gates:     v1.FeatureGates{},
+				ResourceRequirements: internal.ComponentResources{
+					Ingester: internal.ResourceRequirements{
+						PDBMinAvailable: tc.PDBMinAvailable,
+					},
+				},
+				Stack: lokiv1.LokiStackSpec{
+					Template: &lokiv1.LokiTemplateSpec{
+						Ingester: &lokiv1.LokiComponentSpec{
+							Replicas: rand.Int31(),
+						},
+					},
+					Tenants: &lokiv1.TenantsSpec{
+						Mode: lokiv1.OpenshiftLogging,
+					},
+				},
+			}
+			objs, err := manifests.BuildIngester(opts)
+			require.NoError(t, err)
+			require.Len(t, objs, 4)
+
+			pdb := objs[3].(*policyv1.PodDisruptionBudget)
+			require.NotNil(t, pdb)
+			require.Equal(t, "abcd-ingester", pdb.Name)
+			require.Equal(t, "efgh", pdb.Namespace)
+			require.NotNil(t, pdb.Spec.MinAvailable.IntVal)
+			require.Equal(t, int32(tc.ExpectedMinAvailable), pdb.Spec.MinAvailable.IntVal)
+			require.EqualValues(t, manifests.ComponentLabels(manifests.LabelIngesterComponent, opts.Name), pdb.Spec.Selector.MatchLabels)
+		})
+	}
+}

operator/internal/manifests/internal/sizes.go

@@ -1,9 +1,10 @@
 package internal
 
 import (
-	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
+
+	lokiv1 "github.com/grafana/loki/operator/apis/loki/v1"
 )
 
 // ComponentResources is a map of component->requests/limits
@@ -22,9 +23,10 @@ type ComponentResources struct {
 
 // ResourceRequirements sets CPU, Memory, and PVC requirements for a component
 type ResourceRequirements struct {
-	Limits   corev1.ResourceList
-	Requests corev1.ResourceList
-	PVCSize  resource.Quantity
+	Limits          corev1.ResourceList
+	Requests        corev1.ResourceList
+	PVCSize         resource.Quantity
+	PDBMinAvailable int
 }
 
 // ResourceRequirementsTable defines the default resource requests and limits for each size
@@ -66,6 +68,7 @@ var ResourceRequirementsTable = map[lokiv1.LokiStackSizeType]ComponentResources{
 				corev1.ResourceCPU:    resource.MustParse("1"),
 				corev1.ResourceMemory: resource.MustParse("1Gi"),
 			},
+			PDBMinAvailable: 1,
 		},
 		Distributor: corev1.ResourceRequirements{
 			Requests: map[corev1.ResourceName]resource.Quantity{
@@ -123,6 +126,7 @@ var ResourceRequirementsTable = map[lokiv1.LokiStackSizeType]ComponentResources{
 				corev1.ResourceCPU:    resource.MustParse("4"),
 				corev1.ResourceMemory: resource.MustParse("20Gi"),
 			},
+			PDBMinAvailable: 1,
 		},
 		Distributor: corev1.ResourceRequirements{
 			Requests: map[corev1.ResourceName]resource.Quantity{
@@ -180,6 +184,7 @@ var ResourceRequirementsTable = map[lokiv1.LokiStackSizeType]ComponentResources{
 				corev1.ResourceCPU:    resource.MustParse("6"),
 				corev1.ResourceMemory: resource.MustParse("30Gi"),
 			},
+			PDBMinAvailable: 2,
 		},
 		Distributor: corev1.ResourceRequirements{
 			Requests: map[corev1.ResourceName]resource.Quantity{