From afea629a88754315361d600dcfed5552c051d9d1 Mon Sep 17 00:00:00 2001
From: drivebyer <yang.wu@daocloud.io>
Date: Sun, 27 Aug 2023 14:26:25 +0800
Subject: [PATCH] Support redis sentinel pdb

Signed-off-by: drivebyer <yang.wu@daocloud.io>
---
 controllers/redissentinel_controller.go |  5 +++
 example/redis-sentinel.yaml             |  3 ++
 k8sutils/poddisruption.go               | 50 +++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/controllers/redissentinel_controller.go b/controllers/redissentinel_controller.go
index 05b9f5747..84fe332d9 100644
--- a/controllers/redissentinel_controller.go
+++ b/controllers/redissentinel_controller.go
@@ -52,6 +52,11 @@ func (r *RedisSentinelReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 		return ctrl.Result{}, err
 	}
 
+	err = k8sutils.ReconcileSentinelPodDisruptionBudget(instance, instance.Spec.PodDisruptionBudget)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+
 	// Create the Service for Redis Sentinel
 	err = k8sutils.CreateRedisSentinelService(instance)
 	if err != nil {
diff --git a/example/redis-sentinel.yaml b/example/redis-sentinel.yaml
index 4e790b4e2..91b42da9f 100644
--- a/example/redis-sentinel.yaml
+++ b/example/redis-sentinel.yaml
@@ -8,6 +8,9 @@ spec:
   podSecurityContext:
     runAsUser: 1000
     fsGroup: 1000
+  pdb:
+    enabled: false
+    minAvailable: 2
   redisSentinelConfig: 
     redisReplicationName : redis-replication
   kubernetesConfig:
diff --git a/k8sutils/poddisruption.go b/k8sutils/poddisruption.go
index 862f2e9eb..5ef397691 100644
--- a/k8sutils/poddisruption.go
+++ b/k8sutils/poddisruption.go
@@ -38,6 +38,29 @@ func ReconcileRedisPodDisruptionBudget(cr *redisv1beta1.RedisCluster, role strin
 	}
 }
 
+func ReconcileSentinelPodDisruptionBudget(cr *redisv1beta1.RedisSentinel, pdbParams *redisv1beta1.RedisPodDisruptionBudget) error {
+	pdbName := cr.ObjectMeta.Name + "-sentinel"
+	logger := pdbLogger(cr.Namespace, pdbName)
+	if pdbParams != nil && pdbParams.Enabled {
+		labels := getRedisLabels(cr.ObjectMeta.Name, "sentinel", "sentinel", cr.ObjectMeta.GetLabels())
+		annotations := generateStatefulSetsAnots(cr.ObjectMeta)
+		pdbMeta := generateObjectMetaInformation(pdbName, cr.Namespace, labels, annotations)
+		pdbDef := generateSentinelPodDisruptionBudgetDef(cr, "sentinel", pdbMeta, pdbParams)
+		return CreateOrUpdatePodDisruptionBudget(pdbDef)
+	} else {
+		// Check if one exists, and delete it.
+		_, err := GetPodDisruptionBudget(cr.Namespace, pdbName)
+		if err == nil {
+			return deletePodDisruptionBudget(cr.Namespace, pdbName)
+		} else if err != nil && errors.IsNotFound(err) {
+			logger.V(1).Info("Reconciliation Successful, no PodDisruptionBudget Found.")
+			// Its ok if its not found, as we're deleting anyway
+			return nil
+		}
+		return err
+	}
+}
+
 // generatePodDisruptionBudgetDef will create a PodDisruptionBudget definition
 func generatePodDisruptionBudgetDef(cr *redisv1beta1.RedisCluster, role string, pdbMeta metav1.ObjectMeta, pdbParams *redisv1beta1.RedisPodDisruptionBudget) *policyv1.PodDisruptionBudget {
 	lblSelector := LabelSelectors(map[string]string{
@@ -65,6 +88,33 @@ func generatePodDisruptionBudgetDef(cr *redisv1beta1.RedisCluster, role string,
 	return pdbTemplate
 }
 
+// generatePodDisruptionBudgetDef will create a PodDisruptionBudget definition
+func generateSentinelPodDisruptionBudgetDef(cr *redisv1beta1.RedisSentinel, role string, pdbMeta metav1.ObjectMeta, pdbParams *redisv1beta1.RedisPodDisruptionBudget) *policyv1.PodDisruptionBudget {
+	lblSelector := LabelSelectors(map[string]string{
+		"app":  fmt.Sprintf("%s-%s", cr.ObjectMeta.Name, role),
+		"role": role,
+	})
+	pdbTemplate := &policyv1.PodDisruptionBudget{
+		TypeMeta:   generateMetaInformation("PodDisruptionBudget", "policy/v1"),
+		ObjectMeta: pdbMeta,
+		Spec: policyv1.PodDisruptionBudgetSpec{
+			Selector: lblSelector,
+		},
+	}
+	if pdbParams.MinAvailable != nil {
+		pdbTemplate.Spec.MinAvailable = &intstr.IntOrString{Type: intstr.Int, IntVal: int32(*pdbParams.MinAvailable)}
+	}
+	if pdbParams.MaxUnavailable != nil {
+		pdbTemplate.Spec.MaxUnavailable = &intstr.IntOrString{Type: intstr.Int, IntVal: int32(*pdbParams.MaxUnavailable)}
+	}
+	// If we don't have a value for either, assume quorum: (N/2)+1
+	if pdbTemplate.Spec.MaxUnavailable == nil && pdbTemplate.Spec.MinAvailable == nil {
+		pdbTemplate.Spec.MinAvailable = &intstr.IntOrString{Type: intstr.Int, IntVal: int32((*cr.Spec.Size / 2) + 1)}
+	}
+	AddOwnerRefToObject(pdbTemplate, redisSentinelAsOwner(cr))
+	return pdbTemplate
+}
+
 // CreateOrUpdateService method will create or update Redis service
 func CreateOrUpdatePodDisruptionBudget(pdbDef *policyv1.PodDisruptionBudget) error {
 	logger := pdbLogger(pdbDef.Namespace, pdbDef.Name)